Fix up API

posthog/api/test/batch_exports/fixtures.py

@@ -0,0 +1,13 @@
+from posthog.api.test.test_organization import create_organization as create_organization_base
+from posthog.constants import AvailableFeature
+from posthog.models import Organization
+
+
+def create_organization(name: str, has_data_pipelines_feature: bool = True) -> Organization:
+    organization = create_organization_base(name)
+    if has_data_pipelines_feature:
+        organization.available_product_features = [
+            {"key": AvailableFeature.DATA_PIPELINES, "name": AvailableFeature.DATA_PIPELINES}
+        ]
+        organization.save()
+    return organization

posthog/api/test/batch_exports/test_create.py

@@ -9,8 +9,8 @@
 from rest_framework import status
 
 from posthog.api.test.batch_exports.conftest import describe_schedule, start_test_worker
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.batch_exports.operations import create_batch_export
-from posthog.api.test.test_organization import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.batch_exports.models import BatchExport

posthog/api/test/batch_exports/test_delete.py

@@ -8,14 +8,14 @@
 from temporalio.service import RPCError
 
 from posthog.api.test.batch_exports.conftest import start_test_worker
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.batch_exports.operations import (
     backfill_batch_export_ok,
     create_batch_export_ok,
     delete_batch_export,
     delete_batch_export_ok,
     get_batch_export,
 )
-from posthog.api.test.test_organization import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.temporal.common.client import sync_connect

posthog/api/test/batch_exports/test_get.py

@@ -7,7 +7,7 @@
     create_batch_export_ok,
     get_batch_export,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.temporal.common.client import sync_connect
@@ -82,11 +82,11 @@ def test_cannot_get_exports_for_other_organizations(client: HttpClient):
         "interval": "hour",
     }
 
-    organization = create_organization("Test Org")
+    organization = create_organization_with_dat("Test Org")
     team = create_team(organization)
     user = create_user("[email protected]", "Test User", organization)
 
-    another_organization = create_organization("Another Org")
+    another_organization = create_organization_with_dat("Another Org")
     another_user = create_user("[email protected]", "Another Test User", another_organization)
 
     with start_test_worker(temporal):
@@ -126,7 +126,7 @@ def test_batch_exports_are_partitioned_by_team(client: HttpClient):
         "interval": "hour",
     }
 
-    organization = create_organization("Test Org")
+    organization = create_organization_with_dat("Test Org")
     team = create_team(organization)
     another_team = create_team(organization)
     user = create_user("[email protected]", "Test User", organization)

posthog/api/test/batch_exports/test_list.py

@@ -6,7 +6,7 @@
     delete_batch_export_ok,
     list_batch_exports_ok,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 

posthog/api/test/batch_exports/test_pause.py

@@ -14,7 +14,7 @@
     unpause_batch_export,
     unpause_batch_export_ok,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.batch_exports.service import batch_export_delete_schedule

posthog/api/test/batch_exports/test_runs.py

@@ -15,7 +15,7 @@
     get_batch_export_runs,
     get_batch_export_runs_ok,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.temporal.common.client import sync_connect

posthog/api/test/batch_exports/test_update.py

@@ -14,7 +14,7 @@
     patch_batch_export,
     put_batch_export,
 )
-from posthog.api.test.test_organization import create_organization
+from posthog.api.test.batch_exports.fixtures import create_organization
 from posthog.api.test.test_team import create_team
 from posthog.api.test.test_user import create_user
 from posthog.batch_exports.service import sync_batch_export

posthog/api/test/test_team.py

@@ -458,7 +458,10 @@ def test_delete_bulky_postgres_data(self):
         def test_delete_batch_exports(self):
             self.organization_membership.level = OrganizationMembership.Level.ADMIN
             self.organization_membership.save()
-
+            self.organization.available_product_features = [
+                {"key": AvailableFeature.DATA_PIPELINES, "name": AvailableFeature.DATA_PIPELINES}
+            ]
+            self.organization.save()
             team: Team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization)
 
             destination_data = {
@@ -486,16 +489,16 @@ def test_delete_batch_exports(self):
                     json.dumps(batch_export_data),
                     content_type="application/json",
                 )
-                self.assertEqual(response.status_code, 201)
+                assert response.status_code == 201, response.json()
 
                 batch_export = response.json()
                 batch_export_id = batch_export["id"]
 
                 response = self.client.delete(f"/api/environments/{team.id}")
-                self.assertEqual(response.status_code, 204)
+                assert response.status_code == 204, response.json()
 
                 response = self.client.get(f"/api/environments/{team.id}/batch_exports/{batch_export_id}")
-                self.assertEqual(response.status_code, 404)
+                assert response.status_code == 404, response.json()
 
                 with self.assertRaises(RPCError):
                     describe_schedule(temporal, batch_export_id)