Tests: Improve resource teardown by closing/disposing connections etc.

Author: amotl

ingestr/main_test.py

@@ -249,6 +249,8 @@ def test_create_replace_csv_to_duckdb():
     for i, row in enumerate(actual_rows):
         assert res[i] == tuple(row)
 
+    # Clean up
+    conn.close()
     try:
         os.remove(abs_db_path)
     except Exception:
@@ -272,7 +274,10 @@ def test_merge_with_primary_key_csv_to_duckdb():
     rel_db_path_to_command = f"ingestr/testdata/{dbname}"
     uri = f"duckdb:///{rel_db_path_to_command}"
 
-    conn = duckdb.connect(abs_db_path)
+    # DuckDB is sensitive about multiple connections to the same database file.
+    # Connection Error: Can't open a connection to same database file with a
+    # different configuration than existing connections
+    # conn = duckdb.connect(abs_db_path)
 
     def run(source: str):
         res = invoke_ingest_command(
@@ -288,10 +293,13 @@ def run(source: str):
         return res
 
     def get_output_rows():
+        conn = duckdb.connect(abs_db_path)
         conn.execute("CHECKPOINT")
-        return conn.sql(
+        results = conn.sql(
             "select symbol, date, is_enabled, name from testschema_merge.output order by symbol asc"
         ).fetchall()
+        conn.close()
+        return results
 
     def assert_output_equals_to_csv(path: str):
         res = get_output_rows()
@@ -309,19 +317,23 @@ def assert_output_equals_to_csv(path: str):
     run("csv://ingestr/testdata/merge_part1.csv")
     assert_output_equals_to_csv("./testdata/merge_part1.csv")
 
+    conn = duckdb.connect(abs_db_path)
     first_run_id = conn.sql(
         "select _dlt_load_id from testschema_merge.output limit 1"
     ).fetchall()[0][0]
+    conn.close()
 
     ##############################
     # we'll run again, we don't expect any changes since the data hasn't changed
     run("csv://ingestr/testdata/merge_part1.csv")
     assert_output_equals_to_csv("./testdata/merge_part1.csv")
 
     # we also ensure that the other rows were not touched
+    conn = duckdb.connect(abs_db_path)
     count_by_run_id = conn.sql(
         "select _dlt_load_id, count(*) from testschema_merge.output group by 1"
     ).fetchall()
+    conn.close()
     assert len(count_by_run_id) == 1
     assert count_by_run_id[0][1] == 3
     assert count_by_run_id[0][0] == first_run_id
@@ -334,9 +346,11 @@ def assert_output_equals_to_csv(path: str):
     assert_output_equals_to_csv("./testdata/merge_expected.csv")
 
     # let's check the runs
+    conn = duckdb.connect(abs_db_path)
     count_by_run_id = conn.sql(
         "select _dlt_load_id, count(*) from testschema_merge.output group by 1 order by 1 asc"
     ).fetchall()
+    conn.close()
 
     # we expect that there's a new load ID now
     assert len(count_by_run_id) == 2
@@ -349,6 +363,8 @@ def assert_output_equals_to_csv(path: str):
     assert count_by_run_id[1][1] == 3
     ##############################
 
+    # Clean up
+    # conn.close()
     try:
         os.remove(abs_db_path)
     except Exception:
@@ -446,6 +462,8 @@ def assert_output_equals_to_csv(path: str):
     assert count_by_run_id[1][1] == 4
     ##############################
 
+    # Clean up
+    conn.close()
     try:
         os.remove(abs_db_path)
     except Exception:
@@ -786,6 +804,7 @@ def db_to_db_create_replace(
             f"select count(*) from {schema_rand_prefix}.input"
         ).fetchall()
         assert res[0][0] == 2
+    source_engine.dispose()
 
     result = invoke_ingest_command(
         source_connection_url,
@@ -803,6 +822,7 @@ def db_to_db_create_replace(
     res = dest_engine.execute(
         f"select id, val, updated_at from {schema_rand_prefix}.output"
     ).fetchall()
+    dest_engine.dispose()
 
     assert len(res) == 2
 
@@ -845,6 +865,7 @@ def db_to_db_append(
             f"select count(*) from {schema_rand_prefix}.input"
         ).fetchall()
         assert res[0][0] == 2
+    source_engine.dispose()
 
     def run():
         res = invoke_ingest_command(
@@ -858,15 +879,16 @@ def run():
         )
         assert res.exit_code == 0
 
-    dest_engine = sqlalchemy.create_engine(dest_connection_url_read)
-
     def get_output_table():
+        dest_engine = sqlalchemy.create_engine(dest_connection_url_read)
         # CrateDB needs an explicit flush to make data available for reads immediately.
         if dest_engine.dialect.name == "crate":
             dest_engine.execute(f"REFRESH TABLE {schema_rand_prefix}.output")
-        return dest_engine.execute(
+        results = dest_engine.execute(
             f"select id, val, updated_at from {schema_rand_prefix}.output order by id asc"
         ).fetchall()
+        dest_engine.dispose()
+        return results
 
     run()
 
@@ -1011,6 +1033,7 @@ def assert_output_equals(expected):
     source_engine.execute(
         f"UPDATE {schema_rand_prefix}.input SET val = 'val1_modified' WHERE id = 2"
     )
+    source_engine.dispose()
 
     run()
     if dest_connection_url.startswith("cratedb://"):
@@ -1038,6 +1061,7 @@ def assert_output_equals(expected):
     source_engine.execute(
         f"INSERT INTO {schema_rand_prefix}.input VALUES (3, 'val3', '2022-01-01')"
     )
+    source_engine.dispose()
 
     run()
     if dest_connection_url.startswith("cratedb://"):
@@ -1065,6 +1089,7 @@ def assert_output_equals(expected):
     source_engine.execute(
         f"INSERT INTO {schema_rand_prefix}.input VALUES (3, 'val3', '2022-02-02')"
     )
+    source_engine.dispose()
 
     run()
     if dest_connection_url.startswith("cratedb://"):
@@ -1101,6 +1126,7 @@ def assert_output_equals(expected):
     source_engine.execute(
         f"UPDATE {schema_rand_prefix}.input SET val='val2_modified', updated_at = '2022-02-03' WHERE id = 2"
     )
+    source_engine.dispose()
 
     run()
     if dest_connection_url.startswith("cratedb://"):
@@ -1130,6 +1156,7 @@ def assert_output_equals(expected):
     # we don't care about the rest of the run IDs
     assert count_by_run_id[1][1] == 1
     assert count_by_run_id[2][1] == 1
+    dest_engine.dispose()
     ##############################
 
 
@@ -1167,6 +1194,7 @@ def db_to_db_delete_insert_without_primary_key(
             f"select count(*) from {schema_rand_prefix}.input"
         ).fetchall()
         assert res[0][0] == 2
+    source_engine.dispose()
 
     def run():
         res = invoke_ingest_command(
@@ -1189,9 +1217,11 @@ def get_output_rows():
         # CrateDB needs an explicit flush to make data available for reads immediately.
         if dest_engine.dialect.name == "crate":
             dest_engine.execute(f"REFRESH TABLE {schema_rand_prefix}.output")
-        return dest_engine.execute(
+        results = dest_engine.execute(
             f"select id, val, updated_at from {schema_rand_prefix}.output order by id asc"
         ).fetchall()
+        dest_engine.dispose()
+        return results
 
     def assert_output_equals(expected):
         res = get_output_rows()
@@ -1245,6 +1275,7 @@ def assert_output_equals(expected):
     source_engine.execute(
         f"INSERT INTO {schema_rand_prefix}.input VALUES (3, 'val3', '2022-02-01'), (4, 'val4', '2022-02-01')"
     )
+    source_engine.dispose()
 
     run()
     if dest_connection_url.startswith("cratedb://"):
@@ -1314,6 +1345,7 @@ def db_to_db_delete_insert_with_timerange(
         f"select count(*) from {schema_rand_prefix}.input"
     ).fetchall()
     assert res[0][0] == 6
+    source_engine.dispose()
 
     def run(start_date: str, end_date: str):
         res = invoke_ingest_command(
@@ -1451,6 +1483,7 @@ def assert_output_equals(expected):
     source_engine.execute(
         f"UPDATE {schema_rand_prefix}.input SET val = 'val1_modified' WHERE id = 1"
     )
+    source_engine.dispose()
 
     run("2022-01-01", "2022-01-02")
     assert_output_equals(
@@ -1671,6 +1704,7 @@ def run_command(
         ).fetchall()
         assert res[0][0] == "some value"
         assert res[0][1] == row_count
+    dest_engine.dispose()
 
 
 @pytest.mark.parametrize(
@@ -1805,6 +1839,7 @@ def build_datetime(ds: str):
         assert res[0][1] == row_count
         assert res[1][0] == build_datetime("2024-11-06")
         assert res[1][1] == 1000
+    dest_engine.dispose()
 
 
 @pytest.mark.parametrize(
@@ -1915,6 +1950,7 @@ def run_command(df: pd.DataFrame):
         ).fetchall()
         assert res[0][0] == "a"
         assert res[0][1] == row_count + 1000
+    dest_engine.dispose()
 
 
 @pytest.mark.parametrize(
@@ -1952,6 +1988,7 @@ def test_db_to_db_exclude_columns(source, dest):
             f"select count(*) from {schema_rand_prefix}.input"
         ).fetchall()
         assert res[0][0] == 2
+    source_engine.dispose()
     result = invoke_ingest_command(
         source_uri,
         f"{schema_rand_prefix}.input",
@@ -1984,6 +2021,9 @@ def test_db_to_db_exclude_columns(source, dest):
         f"SELECT column_name FROM information_schema.columns WHERE table_schema = '{schema_rand_prefix}' AND table_name = 'output'"
     ).fetchall()
     assert columns == [("id",), ("val",), ("updated_at",)]
+
+    # Clean up
+    dest_engine.dispose()
     source.stop()
     dest.stop()
 
@@ -2404,7 +2444,7 @@ def get_query_result(uri: str, query: str):
     engine = sqlalchemy.create_engine(uri, poolclass=NullPool)
     with engine.connect() as conn:
         res = conn.execute(query).fetchall()
-
+    engine.dispose()
     return res
 
 
@@ -2657,9 +2697,15 @@ def run():
 @pytest.mark.parametrize("source", list(SOURCES.values()), ids=list(SOURCES.keys()))
 @pytest.mark.parametrize("testcase", custom_query_tests())
 def test_custom_query(testcase, source, dest):
-    dest_uri = dest.container.get_connection_url()
+    with ThreadPoolExecutor() as executor:
+        source_future = executor.submit(source.start)
+        dest_future = executor.submit(dest.start)
+        source_uri = source_future.result()
+        dest_uri = dest_future.result()
     dest_uri_read = get_uri_read(dest_uri, dest)
-    testcase(source.start(), dest.start(), dest_uri_read)
+    testcase(source_uri, dest_uri, dest_uri_read)
+    source.stop()
+    dest.stop()
 
 
 # Integration testing when the access token is not provided, and it is only for the resource "repo_events
@@ -2682,6 +2728,7 @@ def test_github_to_duckdb(dest):
     if dest_engine.dialect.name == "crate":
         dest_engine.execute(f"REFRESH TABLE {dest_table}")
     res = dest_engine.execute(f"select count(*) from {dest_table}").fetchall()
+    dest_engine.dispose()
     assert len(res) > 0
 
 
@@ -2967,6 +3014,7 @@ def test_successful_ingestion(dest_uri, dest_uri_read):
         if dest_engine.dialect.name == "crate":
             dest_engine.execute(f"REFRESH TABLE {dest_table}")
         count = dest_engine.execute(f"select count(*) from {dest_table}").fetchone()[0]
+        dest_engine.dispose()
         assert count == 3
 
     def test_incremental_ingestion(dest_uri, dest_uri_read):
@@ -3115,6 +3163,7 @@ def test_incremental_ingestion(dest_uri, dest_uri_read):
             )
             == 2
         )
+        dest_engine.dispose()
 
     return [
         test_no_report_instances_found,
@@ -3209,6 +3258,7 @@ def assert_rows(dest_uri, dest_table, n):
             rows = conn.execute(f"select count(*) from {dest_table}").fetchall()
             assert len(rows) == 1
             assert rows[0] == (n,)
+        engine.dispose()
 
     def test_empty_source_uri(dest_uri, dest_uri_read):
         """
@@ -3628,6 +3678,7 @@ def exchange_rate_on_specific_date(dest_uri, dest_uri_read):
             if dest_engine.dialect.name == "crate":
                 conn.execute(f"REFRESH TABLE {dest_table}")
             rows = conn.execute(query).fetchall()
+        dest_engine.dispose()
 
         # Assert that the rate for GBP is 0.82993
         assert len(rows) > 0, "No data found for GBP"
@@ -3710,6 +3761,7 @@ def test_mysql_zero_dates(source, dest):
             f"select count(*) from {schema_rand_prefix}.input"
         ).fetchall()
         assert res[0][0] == 5
+    source_engine.dispose()
 
     result = invoke_ingest_command(
         source_uri,
@@ -3728,6 +3780,7 @@ def test_mysql_zero_dates(source, dest):
     if dest_engine.dialect.name == "crate":
         dest_engine.execute(f"REFRESH TABLE {schema_rand_prefix}.output")
     res = dest_engine.execute(f"select * from {schema_rand_prefix}.output").fetchall()
+    dest_engine.dispose()
 
     # assert there are no new rows, since DBs like DuckDB accept NULL and dlt adds a separate string column for the value `0000-00-00 00:00:00`
     # we want 4 columns: name, created_at, _dlt_load_id, _dlt_id
@@ -4197,6 +4250,7 @@ def test_stripe_source_full_refresh(stripe_table):
     assert res[0] > 0, f"No {stripe_table} records found"
 
     # Clean up
+    conn.close()
     try:
         os.remove(abs_db_path)
     except Exception:
@@ -4237,6 +4291,7 @@ def test_stripe_source_incremental(stripe_table):
     assert res[0] > 0, f"No {stripe_table} records found"
 
     # Clean up
+    conn.close()
     try:
         os.remove(abs_db_path)
     except Exception: