update test

PostHog · Dec 26, 2023 · c33fe1a · c33fe1a
1 parent 1188f62
commit c33fe1a
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/posthog/temporal/tests/test_external_data_job.py b/posthog/temporal/tests/test_external_data_job.py
@@ -22,6 +22,7 @@
 )
 from posthog.warehouse.models import (
     get_latest_run_if_exists,
+    get_all_schemas_for_source_id,
     DataWarehouseTable,
     ExternalDataJob,
     ExternalDataSource,
@@ -30,6 +31,7 @@
 
 from posthog.temporal.data_imports.pipelines.schemas import (
     PIPELINE_TYPE_SCHEMA_DEFAULT_MAPPING,
+    PIPELINE_TYPE_SCHEMA_MAPPING,
 )
 from posthog.temporal.data_imports.pipelines.pipeline import DataImportPipeline
 from temporalio.testing import WorkflowEnvironment
@@ -112,6 +114,9 @@ async def test_create_external_job_activity(activity_environment, team, **kwargs
     assert await sync_to_async(runs.exists)()  # type:ignore
     assert len(schemas) == len(PIPELINE_TYPE_SCHEMA_DEFAULT_MAPPING[new_source.source_type])
 
+    all_schemas = await sync_to_async(get_all_schemas_for_source_id)(new_source.pk, team_id=team.pk)
+    assert len(all_schemas) == len(PIPELINE_TYPE_SCHEMA_MAPPING[new_source.source_type])
+
 
 @pytest.mark.django_db(transaction=True)
 @pytest.mark.asyncio