From e0bc92b1c87474f6118fff26a0eb04803a66e342 Mon Sep 17 00:00:00 2001 From: Eric Duong Date: Fri, 21 Jun 2024 13:54:03 -0400 Subject: [PATCH] fix(data-warehouse): test prefetch (#23162) make sure job prefetches --- .../external_data/test_external_data_job.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 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 index 01224cb58c410..bbaccf7308f9b 100644 --- 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