Skip to content

Commit

Permalink
Fixed test_python_query_handler_runner.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@ckunki
ckunki committed Nov 7, 2024
1 parent fec6805 commit 3edc746
Showing 1 changed file with 60 additions and 38 deletions.
98 changes: 60 additions & 38 deletions tests/unit_tests/query_handler_runner/test_python_query_handler_runner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,14 +30,24 @@ def expect_query(template: str, result_set = MockResultSet()):
return [template, result_set]


def create_sql_executor(schema: str, *args):
def create_sql_executor_1(schema: str, *args):
return MockSQLExecutor([
ExpectedQuery(
cleandoc(template.format(schema=schema)),
result_set or MockResultSet()
) for template, result_set in args
])


def create_sql_executor(schema: str, prefix: str, *args):
return MockSQLExecutor([
ExpectedQuery(
cleandoc(template.format(schema=schema, prefix=prefix)),
result_set or MockResultSet()
) for template, result_set in args
])


@pytest.fixture()
def prefix(tmp_db_obj_prefix):
return tmp_db_obj_prefix
Expand Down Expand Up @@ -107,8 +117,13 @@ def test_start_finish_cleanup_queries(aaf_pytest_db_schema, prefix, context_mock
and then directly returns a Finish result. We expect a cleanup query for the temporary
table to be executed and result of the Finish object returned.
"""
expected = f"""DROP TABLE IF EXISTS "{temporary_schema_name}"."temp_db_object_1";""";
sql_executor = MockSQLExecutor([ExpectedQuery(expected)])
sql_executor = create_sql_executor(
aaf_pytest_db_schema,
prefix,
expect_query(
'DROP TABLE IF EXISTS "{schema}"."{prefix}_1";'
))

test_input = TestInput()
query_handler_runner = PythonQueryHandlerRunner[TestInput, TestOutput](
sql_executor=sql_executor,
Expand Down Expand Up @@ -140,9 +155,10 @@ def test_start_error_cleanup_queries(aaf_pytest_db_schema, prefix, context_mock)
table to be executed and the exception to be forwarded.
"""
sql_executor = create_sql_executor(
temporary_schema_name,
aaf_pytest_db_schema,
prefix,
expect_query(
'DROP TABLE IF EXISTS "{schema}"."temp_db_object_1";'
'DROP TABLE IF EXISTS "{schema}"."{prefix}_1";'
))

test_input = TestInput()
Expand Down Expand Up @@ -185,26 +201,27 @@ def test_continue_finish(aaf_pytest_db_schema, prefix, context_mock):
in the input_query.
"""
sql_executor = create_sql_executor(
temporary_schema_name,
aaf_pytest_db_schema,
prefix,
expect_query(
"""
CREATE OR REPLACE VIEW "{schema}"."temp_db_object_2_1" AS
CREATE OR REPLACE VIEW "{schema}"."{prefix}_2_1" AS
SELECT 1 as "a";
"""
),
expect_query(
"""
SELECT
"a"
FROM "{schema}"."temp_db_object_2_1";
FROM "{schema}"."{prefix}_2_1";
""",
MockResultSet(
rows=[(1,)],
columns=[Column(ColumnName("a"), ColumnType(
name="DECIMAL", precision=1, scale=0))])
),
expect_query(
'DROP VIEW IF EXISTS "{schema}"."temp_db_object_2_1";'
'DROP VIEW IF EXISTS "{schema}"."{prefix}_2_1";'
),
)

Expand Down Expand Up @@ -233,34 +250,35 @@ def handle_query_result(self, query_result: QueryResult) -> Union[Continue, Fini
raise AssertionError("handle_query_result shouldn't be called")


def test_continue_wrong_columns(aaf_pytest_db_schema, context_mock):
def test_continue_wrong_columns(aaf_pytest_db_schema, prefix, context_mock):
"""
This tests runs a query handler which returns Continue result with mismatching column definition
between the input query and its column definition. We expect the query handler runner to raise
an Exception.
"""

sql_executor = create_sql_executor(
temporary_schema_name,
aaf_pytest_db_schema,
prefix,
expect_query(
"""
CREATE OR REPLACE VIEW "{schema}"."temp_db_object_2_1" AS
CREATE OR REPLACE VIEW "{schema}"."{prefix}_2_1" AS
SELECT 1 as "b";
"""
),
expect_query(
"""
SELECT
"a"
FROM "{schema}"."temp_db_object_2_1";
FROM "{schema}"."{prefix}_2_1";
""",
MockResultSet(
rows=[(1,)],
columns=[Column(ColumnName("b"), ColumnType(
name="DECIMAL", precision=1, scale=0))])
),
expect_query(
'DROP VIEW IF EXISTS "{schema}"."temp_db_object_2_1";'
'DROP VIEW IF EXISTS "{schema}"."{prefix}_2_1";'
),
)
test_input = TestInput()
Expand Down Expand Up @@ -302,11 +320,12 @@ def test_continue_query_list(aaf_pytest_db_schema, prefix, context_mock):
"""

sql_executor = create_sql_executor(
temporary_schema_name,
aaf_pytest_db_schema,
prefix,
expect_query("SELECT 1"),
expect_query(
"""
CREATE OR REPLACE VIEW "{schema}"."temp_db_object_2_1" AS
CREATE OR REPLACE VIEW "{schema}"."{prefix}_2_1" AS
SELECT 1 as "a";
""",
MockResultSet(
Expand All @@ -318,15 +337,15 @@ def test_continue_query_list(aaf_pytest_db_schema, prefix, context_mock):
"""
SELECT
"a"
FROM "{schema}"."temp_db_object_2_1";
FROM "{schema}"."{prefix}_2_1";
""",
MockResultSet(
rows=[(1,)],
columns=[Column(ColumnName("a"), ColumnType(
name="DECIMAL", precision=1, scale=0))])
),
expect_query(
'DROP VIEW IF EXISTS "{schema}"."temp_db_object_2_1";'
'DROP VIEW IF EXISTS "{schema}"."{prefix}_2_1";'
),
)
test_input = TestInput()
Expand Down Expand Up @@ -367,25 +386,26 @@ def test_continue_error_cleanup_queries(aaf_pytest_db_schema, prefix, context_mo
"""

sql_executor = create_sql_executor(
temporary_schema_name,
aaf_pytest_db_schema,
prefix,
expect_query(
"""
CREATE OR REPLACE VIEW "{schema}"."temp_db_object_2_1" AS
CREATE OR REPLACE VIEW "{schema}"."{prefix}_2_1" AS
SELECT 1 as "a";
"""),
expect_query(
"""
SELECT
"a"
FROM "{schema}"."temp_db_object_2_1";
FROM "{schema}"."{prefix}_2_1";
""",
MockResultSet(
rows=[(1,)],
columns=[Column(ColumnName("a"), ColumnType(
name="DECIMAL", precision=1, scale=0))])
),
expect_query('DROP TABLE IF EXISTS "{schema}"."temp_db_object_3";'),
expect_query('DROP VIEW IF EXISTS "{schema}"."temp_db_object_2_1";'),
expect_query('DROP TABLE IF EXISTS "{schema}"."{prefix}_3";'),
expect_query('DROP VIEW IF EXISTS "{schema}"."{prefix}_2_1";'),
)
test_input = TestInput()
query_handler_runner = PythonQueryHandlerRunner[TestInput, TestOutput](
Expand Down Expand Up @@ -438,44 +458,45 @@ def test_continue_continue_finish(aaf_pytest_db_schema, prefix, context_mock):
"""

sql_executor = create_sql_executor(
temporary_schema_name,
aaf_pytest_db_schema,
prefix,
expect_query(
"""
CREATE OR REPLACE VIEW "{schema}"."temp_db_object_2_1" AS
CREATE OR REPLACE VIEW "{schema}"."{prefix}_2_1" AS
SELECT 1 as "a";
"""),
expect_query(
"""
SELECT
"a"
FROM "{schema}"."temp_db_object_2_1";
FROM "{schema}"."{prefix}_2_1";
""",
MockResultSet(
rows=[(1,)],
columns=[Column(ColumnName("a"), ColumnType(
name="DECIMAL", precision=1, scale=0))])
),
expect_query(
'DROP VIEW IF EXISTS "{schema}"."temp_db_object_2_1";'
'DROP VIEW IF EXISTS "{schema}"."{prefix}_2_1";'
),
expect_query(
"""
CREATE OR REPLACE VIEW "{schema}"."temp_db_object_4_1" AS
CREATE OR REPLACE VIEW "{schema}"."{prefix}_4_1" AS
SELECT 1 as "b";
"""),
expect_query(
"""
SELECT
"b"
FROM "{schema}"."temp_db_object_4_1";
FROM "{schema}"."{prefix}_4_1";
""",
MockResultSet(
rows=[(1,)],
columns=[Column(ColumnName("b"), ColumnType(
name="DECIMAL", precision=1, scale=0))])
),
expect_query(
'DROP VIEW IF EXISTS "{schema}"."temp_db_object_4_1";'),
'DROP VIEW IF EXISTS "{schema}"."{prefix}_4_1";'),
)
test_input = TestInput()
query_handler_runner = PythonQueryHandlerRunner[TestInput, TestOutput](
Expand Down Expand Up @@ -530,47 +551,48 @@ def test_continue_cleanup_continue_finish(aaf_pytest_db_schema, prefix, context_
"""

sql_executor = create_sql_executor(
temporary_schema_name,
aaf_pytest_db_schema,
prefix,
expect_query(
"""
CREATE OR REPLACE VIEW "{schema}"."temp_db_object_4_1" AS
CREATE OR REPLACE VIEW "{schema}"."{prefix}_4_1" AS
SELECT 1 as "a";
"""),
expect_query(
"""
SELECT
"a"
FROM "{schema}"."temp_db_object_4_1";
FROM "{schema}"."{prefix}_4_1";
""",
MockResultSet(
rows=[(1,)],
columns=[Column(ColumnName("a"), ColumnType(
name="DECIMAL", precision=1, scale=0))])
),
expect_query(
'DROP VIEW IF EXISTS "{schema}"."temp_db_object_4_1";'
'DROP VIEW IF EXISTS "{schema}"."{prefix}_4_1";'
),
expect_query(
'DROP TABLE IF EXISTS "{schema}"."temp_db_object_2_1";'
'DROP TABLE IF EXISTS "{schema}"."{prefix}_2_1";'
),
expect_query(
"""
CREATE OR REPLACE VIEW "{schema}"."temp_db_object_6_1" AS
CREATE OR REPLACE VIEW "{schema}"."{prefix}_6_1" AS
SELECT 1 as "b";
"""),
expect_query(
"""
SELECT
"b"
FROM "{schema}"."temp_db_object_6_1";
FROM "{schema}"."{prefix}_6_1";
""",
MockResultSet(
rows=[(1,)],
columns=[Column(ColumnName("b"), ColumnType(
name="DECIMAL", precision=1, scale=0))])
),
expect_query(
'DROP VIEW IF EXISTS "{schema}"."temp_db_object_6_1";'
'DROP VIEW IF EXISTS "{schema}"."{prefix}_6_1";'
),
)
test_input = TestInput()
Expand Down

0 comments on commit 3edc746

Please sign in to comment.