feat: add wizard e2e test query edge case

.run/PostHog.run.xml

@@ -1,10 +1,11 @@
 <component name="ProjectRunConfigurationManager">
   <configuration default="false" name="PostHog" type="Python.DjangoServer" factoryName="Django server">
     <module name="posthog" />
-    <option name="ENV_FILES" value="" />
+    <option name="ENV_FILES" value="$PROJECT_DIR$/.env.local" />
     <option name="INTERPRETER_OPTIONS" value="" />
     <option name="PARENT_ENVS" value="true" />
     <envs>
+      <env name="PYTHONUNBUFFERED" value="1" />
       <env name="BILLING_SERVICE_URL" value="https://billing.dev.posthog.dev" />
       <env name="CAPTURE_TIME_TO_SEE_DATA" value="0" />
       <env name="CLICKHOUSE_SECURE" value="False" />
@@ -15,7 +16,6 @@
       <env name="KEA_VERBOSE_LOGGING" value="false" />
       <env name="PRINT_SQL" value="1" />
       <env name="PYDEVD_USE_CYTHON" value="NO" />
-      <env name="PYTHONUNBUFFERED" value="1" />
       <env name="SESSION_RECORDING_KAFKA_COMPRESSION" value="gzip-in-capture" />
       <env name="SESSION_RECORDING_KAFKA_HOSTS" value="localhost" />
       <env name="SESSION_RECORDING_KAFKA_MAX_REQUEST_SIZE_BYTES" value="20971520" />

posthog/api/wizard/http.py

@@ -148,13 +148,25 @@ def query(self, request: Request) -> Response:
         model = validated_data["model"]
 
         hash = request.headers.get("X-PostHog-Wizard-Hash")
+        fixture_generation = request.headers.get("X-PostHog-Wizard-Fixture-Generation")
 
         if not hash:
             raise AuthenticationFailed("X-PostHog-Wizard-Hash header is required.")
 
         key = f"{SETUP_WIZARD_CACHE_PREFIX}{hash}"
         wizard_data = cache.get(key)
 
+        # wizard_data should only be mocked during the @posthog/wizard E2E tests, so that fixtures can be generated.
+        mock_wizard_data = settings.DEBUG and fixture_generation
+
+        if mock_wizard_data:
+            wizard_data = {
+                "project_api_key": "mock-project-api-key",
+                "host": "http://localhost:8010",
+                "user_distinct_id": "mock-user-id",
+            }
+            cache.set(key, wizard_data, SETUP_WIZARD_CACHE_TIMEOUT)
+
         if wizard_data is None:
             raise AuthenticationFailed("Invalid hash.")
 
@@ -232,6 +244,7 @@ def query(self, request: Request) -> Response:
                     "ai_product": "wizard",
                     "ai_feature": "query",
                 },
+                temperature=1.0,
             )
 
             if (

posthog/api/wizard/test/test_http.py

@@ -52,6 +52,7 @@ def test_query_endpoint_requires_hash_header(self, mock_openai):
 
     @patch("posthog.api.wizard.http.posthoganalytics.default_client", MagicMock())
     @patch("posthog.api.wizard.http.OpenAI")
+    @patch("django.conf.settings.DEBUG", False)
     def test_query_endpoint_rate_limit(self, mock_openai):
         mock_openai_instance = mock_openai.return_value
         # Simulate an OpenAI response with JSON {"foo": "bar"}
@@ -206,6 +207,126 @@ def test_query_endpoint_rejects_invalid_model(self):
         assert "model" in response.json()
         assert "not supported" in response.json()["model"][0]
 
+    @patch("django.conf.settings.DEBUG", True)
+    @patch("posthog.api.wizard.http.posthoganalytics.default_client", MagicMock())
+    @patch("posthog.api.wizard.http.OpenAI")
+    def test_query_endpoint_mock_wizard_data_in_debug_with_fixture_header(self, mock_openai):
+        """Test that mock wizard data is used when DEBUG=True and X-PostHog-Wizard-Fixture-Generation header is present"""
+        mock_openai_instance = mock_openai.return_value
+        mock_openai_instance.chat.completions.create.return_value = MagicMock(
+            choices=[MagicMock(message=MagicMock(content=json.dumps({"result": "mocked"})))]
+        )
+
+        # Clear any existing cache data
+        cache.delete(self.cache_key)
+
+        response = self.client.post(
+            self.query_url,
+            data=json.dumps(
+                {
+                    "message": "test",
+                    "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}},
+                }
+            ),
+            content_type="application/json",
+            HTTP_X_POSTHOG_WIZARD_HASH=self.hash,
+            HTTP_X_POSTHOG_WIZARD_FIXTURE_GENERATION="true",
+        )
+
+        assert response.status_code == status.HTTP_200_OK
+        assert response.json() == {"data": {"result": "mocked"}}
+
+        # Verify that mock data was cached
+        cached_data = cache.get(self.cache_key)
+        assert cached_data is not None
+        assert cached_data["project_api_key"] == "mock-project-api-key"
+        assert cached_data["host"] == "http://localhost:8010"
+        assert cached_data["user_distinct_id"] == "mock-user-id"
+
+    @patch("django.conf.settings.DEBUG", True)
+    @patch("posthog.api.wizard.http.posthoganalytics.default_client", MagicMock())
+    @patch("posthog.api.wizard.http.OpenAI")
+    def test_query_endpoint_mock_wizard_data_overrides_existing_cache(self, mock_openai):
+        """Test that mock wizard data overrides existing cache data when conditions are met"""
+        mock_openai_instance = mock_openai.return_value
+        mock_openai_instance.chat.completions.create.return_value = MagicMock(
+            choices=[MagicMock(message=MagicMock(content=json.dumps({"result": "overridden"})))]
+        )
+
+        # Set existing cache data
+        cache.set(
+            self.cache_key, {"project_api_key": "real-key", "host": "https://real-host.com"}, SETUP_WIZARD_CACHE_TIMEOUT
+        )
+
+        response = self.client.post(
+            self.query_url,
+            data=json.dumps(
+                {
+                    "message": "test",
+                    "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}},
+                }
+            ),
+            content_type="application/json",
+            HTTP_X_POSTHOG_WIZARD_HASH=self.hash,
+            HTTP_X_POSTHOG_WIZARD_FIXTURE_GENERATION="true",
+        )
+
+        assert response.status_code == status.HTTP_200_OK
+        assert response.json() == {"data": {"result": "overridden"}}
+
+        # Verify that cache was overridden with mock data
+        cached_data = cache.get(self.cache_key)
+        assert cached_data["project_api_key"] == "mock-project-api-key"
+        assert cached_data["host"] == "http://localhost:8010"
+        assert cached_data["user_distinct_id"] == "mock-user-id"
+
+    @patch("django.conf.settings.DEBUG", False)
+    @patch("posthog.api.wizard.http.posthoganalytics.default_client", MagicMock())
+    @patch("posthog.api.wizard.http.OpenAI")
+    def test_query_endpoint_no_mock_when_debug_false(self, mock_openai):
+        """Test that mock wizard data is NOT used when DEBUG=False even with fixture header"""
+        # Clear any existing cache data
+        cache.delete(self.cache_key)
+
+        response = self.client.post(
+            self.query_url,
+            data=json.dumps(
+                {
+                    "message": "test",
+                    "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}},
+                }
+            ),
+            content_type="application/json",
+            HTTP_X_POSTHOG_WIZARD_HASH=self.hash,
+            HTTP_X_POSTHOG_WIZARD_FIXTURE_GENERATION="true",
+        )
+
+        # Should fail authentication because no cache data exists and mock is not used
+        assert response.status_code == status.HTTP_401_UNAUTHORIZED
+
+    @patch("django.conf.settings.DEBUG", True)
+    @patch("posthog.api.wizard.http.posthoganalytics.default_client", MagicMock())
+    @patch("posthog.api.wizard.http.OpenAI")
+    def test_query_endpoint_no_mock_without_fixture_header(self, mock_openai):
+        """Test that mock wizard data is NOT used when DEBUG=True but fixture header is missing"""
+        # Clear any existing cache data
+        cache.delete(self.cache_key)
+
+        response = self.client.post(
+            self.query_url,
+            data=json.dumps(
+                {
+                    "message": "test",
+                    "json_schema": {"type": "object", "properties": {"name": {"type": "string"}}},
+                }
+            ),
+            content_type="application/json",
+            HTTP_X_POSTHOG_WIZARD_HASH=self.hash,
+        )
+
+        # Should fail authentication because no cache data exists and mock is not used
+        assert response.status_code == status.HTTP_401_UNAUTHORIZED
+
     @override_settings(
         CACHES={
             "default": {

posthog/rate_limit.py

@@ -3,7 +3,7 @@
 import time
 from functools import lru_cache
 from typing import Optional
-
+from django.conf import settings
 from prometheus_client import Counter
 from rest_framework.throttling import SimpleRateThrottle, BaseThrottle, UserRateThrottle
 from rest_framework.request import Request
@@ -425,7 +425,10 @@ class SetupWizardAuthenticationRateThrottle(UserRateThrottle):
 
 
 class SetupWizardQueryRateThrottle(SimpleRateThrottle):
-    rate = "20/day"  # Since the authentication hash is valid for a short period, this is effectively per-user
+    def get_rate(self):
+        if settings.DEBUG:
+            return "1000/day"
+        return "20/day"
 
     # Throttle per wizard hash
     def get_cache_key(self, request, view):