Resolve sqlite deprecations in tests (apache#40389)

astronomer · Jun 23, 2024 · fd44198 · fd44198
1 parent 68a87f7
commit fd44198
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/tests/deprecations_ignore.yml b/tests/deprecations_ignore.yml
@@ -248,7 +248,4 @@
 - tests/providers/snowflake/operators/test_snowflake.py::TestSnowflakeOperatorForParams::test_overwrite_params
 - tests/providers/snowflake/operators/test_snowflake_sql.py::test_exec_success
 - tests/providers/snowflake/operators/test_snowflake_sql.py::test_execute_openlineage_events
-- tests/providers/sqlite/operators/test_sqlite.py::TestSqliteOperator::test_sqlite_operator_with_invalid_sql
-- tests/providers/sqlite/operators/test_sqlite.py::TestSqliteOperator::test_sqlite_operator_with_multiple_statements
-- tests/providers/sqlite/operators/test_sqlite.py::TestSqliteOperator::test_sqlite_operator_with_one_statement
 - tests/providers/trino/operators/test_trino.py::test_execute_openlineage_events
diff --git a/tests/providers/sqlite/operators/test_sqlite.py b/tests/providers/sqlite/operators/test_sqlite.py
@@ -20,7 +20,7 @@
 import pytest
 
 from airflow.models.dag import DAG
-from airflow.providers.sqlite.operators.sqlite import SqliteOperator
+from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
 from airflow.utils import timezone
 
 DEFAULT_DATE = timezone.datetime(2015, 1, 1)
@@ -51,15 +51,20 @@ def test_sqlite_operator_with_one_statement(self):
             dummy VARCHAR(50)
         );
         """
-        op = SqliteOperator(task_id="basic_sqlite", sql=sql, dag=self.dag)
+        op = SQLExecuteQueryOperator(task_id="basic_sqlite", sql=sql, dag=self.dag, conn_id="sqlite_default")
         op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
 
     def test_sqlite_operator_with_multiple_statements(self):
         sql = [
             "CREATE TABLE IF NOT EXISTS test_airflow (dummy VARCHAR(50))",
             "INSERT INTO test_airflow VALUES ('X')",
         ]
-        op = SqliteOperator(task_id="sqlite_operator_with_multiple_statements", sql=sql, dag=self.dag)
+        op = SQLExecuteQueryOperator(
+            task_id="sqlite_operator_with_multiple_statements",
+            sql=sql,
+            dag=self.dag,
+            conn_id="sqlite_default",
+        )
         op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
 
     def test_sqlite_operator_with_invalid_sql(self):
@@ -70,6 +75,11 @@ def test_sqlite_operator_with_invalid_sql(self):
 
         from sqlite3 import OperationalError
 
-        op = SqliteOperator(task_id="sqlite_operator_with_multiple_statements", sql=sql, dag=self.dag)
+        op = SQLExecuteQueryOperator(
+            task_id="sqlite_operator_with_multiple_statements",
+            sql=sql,
+            dag=self.dag,
+            conn_id="sqlite_default",
+        )
         with pytest.raises(OperationalError, match="no such table: test_airflow2"):
             op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)