additional conditions for assert

dlt-hub · May 7, 2024 · 64eaf88 · 64eaf88
1 parent 97a9033
commit 64eaf88
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/tests/load/pipeline/test_pipelines.py b/tests/load/pipeline/test_pipelines.py
@@ -996,12 +996,14 @@ def table_3(make_data=False):
         is True
     )
 
-    client, _ = pipeline._get_destination_clients(schema)
+    job_client, _ = pipeline._get_destination_clients(schema)
 
-    if isinstance(client, WithStagingDataset):
+    if isinstance(job_client, WithStagingDataset):
         with pipeline.sql_client() as client:
             for i in range(1, 4):
-                if client.should_load_data_to_staging_dataset(f"table_{i}"):
+                if job_client.should_load_data_to_staging_dataset(
+                    pipeline.default_schema.tables[f"table_{i}"]
+                ):
                     with client.execute_query(
                         f"SELECT * FROM {pipeline.dataset_name}_staging.table_{i}"
                     ) as cur: