fix test

PostHog · Oct 17, 2024 · 4422580 · 4422580
1 parent b2d9e80
commit 4422580
Showing 1 changed file with 6 additions and 6 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
@@ -27,6 +27,7 @@
     ExternalDataJob,
     ExternalDataSource,
     ExternalDataSchema,
+    get_external_data_job,
 )
 
 from posthog.temporal.data_imports.pipelines.schemas import (
@@ -379,9 +380,7 @@ async def setup_job_1():
             schema=customer_schema,
         )
 
-        new_job = await sync_to_async(
-            ExternalDataJob.objects.filter(id=new_job.id).prefetch_related("pipeline").prefetch_related("schema").get
-        )()
+        new_job = await get_external_data_job(new_job.id)
 
         inputs = ImportDataActivityInputs(
             team_id=team.id,
@@ -403,16 +402,17 @@ async def setup_job_2():
             job_inputs={"stripe_secret_key": "test-key", "stripe_account_id": "acct_id"},
         )
 
+        charge_schema = await _create_schema("Charge", new_source, team)
+
         new_job: ExternalDataJob = await sync_to_async(ExternalDataJob.objects.create)(
             team_id=team.id,
             pipeline_id=new_source.pk,
             status=ExternalDataJob.Status.RUNNING,
             rows_synced=0,
+            schema=charge_schema,
         )
 
-        new_job = await sync_to_async(ExternalDataJob.objects.filter(id=new_job.id).prefetch_related("pipeline").get)()
-
-        charge_schema = await _create_schema("Charge", new_source, team)
+        new_job = await get_external_data_job(new_job.id)
 
         inputs = ImportDataActivityInputs(
             team_id=team.id,