fix(data-warehouse): test prefetch (#23162)

make sure job prefetches
PostHog · Jun 21, 2024 · e0bc92b · e0bc92b
1 parent 186711f
commit e0bc92b
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/posthog/temporal/tests/external_data/test_external_data_job.py b/posthog/temporal/tests/external_data/test_external_data_job.py
@@ -283,7 +283,9 @@ async def setup_job_1():
             schema=customer_schema,
         )
 
-        new_job = await sync_to_async(ExternalDataJob.objects.filter(id=new_job.id).prefetch_related("pipeline").get)()
+        new_job = await sync_to_async(
+            ExternalDataJob.objects.filter(id=new_job.id).prefetch_related("pipeline").prefetch_related("schema").get
+        )()
 
         inputs = ImportDataActivityInputs(
             team_id=team.id,
@@ -438,7 +440,9 @@ async def setup_job_1():
             schema=customer_schema,
         )
 
-        new_job = await sync_to_async(ExternalDataJob.objects.filter(id=new_job.id).prefetch_related("pipeline").get)()
+        new_job = await sync_to_async(
+            ExternalDataJob.objects.filter(id=new_job.id).prefetch_related("pipeline").prefetch_related("schema").get
+        )()
 
         inputs = ImportDataActivityInputs(
             team_id=team.id,
@@ -522,7 +526,9 @@ async def setup_job_1():
             schema=customer_schema,
         )
 
-        new_job = await sync_to_async(ExternalDataJob.objects.filter(id=new_job.id).prefetch_related("pipeline").get)()
+        new_job = await sync_to_async(
+            ExternalDataJob.objects.filter(id=new_job.id).prefetch_related("pipeline").prefetch_related("schema").get
+        )()
 
         inputs = ImportDataActivityInputs(
             team_id=team.id,
@@ -688,7 +694,9 @@ async def setup_job_1():
             schema=posthog_test_schema,
         )
 
-        new_job = await sync_to_async(ExternalDataJob.objects.filter(id=new_job.id).prefetch_related("pipeline").get)()
+        new_job = await sync_to_async(
+            ExternalDataJob.objects.filter(id=new_job.id).prefetch_related("pipeline").prefetch_related("schema").get
+        )()
 
         inputs = ImportDataActivityInputs(
             team_id=team.id, run_id=str(new_job.pk), source_id=new_source.pk, schema_id=posthog_test_schema.id