From f5f365080bfdcc9a3b261637c1bb7c16a156a122 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 21 Jul 2025 18:07:16 +0200 Subject: [PATCH 1/4] feat: add wizard e2e test query edge case --- posthog/api/wizard/http.py | 13 +++ posthog/api/wizard/test/test_http.py | 120 +++++++++++++++++++++++++++ posthog/rate_limit.py | 5 +- 3 files changed, 137 insertions(+), 1 deletion(-) diff --git a/posthog/api/wizard/http.py b/posthog/api/wizard/http.py index e975a854ba17..9bec87af8441 100644 --- a/posthog/api/wizard/http.py +++ b/posthog/api/wizard/http.py @@ -130,6 +130,7 @@ 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.") @@ -137,6 +138,17 @@ def query(self, request: Request) -> Response: 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.") @@ -214,6 +226,7 @@ def query(self, request: Request) -> Response: "ai_product": "wizard", "ai_feature": "query", }, + temperature=1.0, ) if ( diff --git a/posthog/api/wizard/test/test_http.py b/posthog/api/wizard/test/test_http.py index bda108512fba..ce5f3d6c018f 100644 --- a/posthog/api/wizard/test/test_http.py +++ b/posthog/api/wizard/test/test_http.py @@ -202,6 +202,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 + def tearDown(self): super().tearDown() cache.clear() # Clears out all DRF throttle data diff --git a/posthog/rate_limit.py b/posthog/rate_limit.py index 5d5920f327e1..1e4e0a722379 100644 --- a/posthog/rate_limit.py +++ b/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 @@ -427,6 +427,9 @@ class SetupWizardAuthenticationRateThrottle(UserRateThrottle): class SetupWizardQueryRateThrottle(SimpleRateThrottle): rate = "20/day" # Since the authentication hash is valid for a short period, this is effectively per-user + if settings.DEBUG: + "1000/day" + # Throttle per wizard hash def get_cache_key(self, request, view): hash = request.headers.get("X-PostHog-Wizard-Hash") From 6aecf9d0da1239e0923d8678793d774e48bf7a81 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 21 Jul 2025 18:14:22 +0200 Subject: [PATCH 2/4] typo fix --- posthog/rate_limit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/rate_limit.py b/posthog/rate_limit.py index 1e4e0a722379..392ecaba8bd0 100644 --- a/posthog/rate_limit.py +++ b/posthog/rate_limit.py @@ -428,7 +428,7 @@ class SetupWizardQueryRateThrottle(SimpleRateThrottle): rate = "20/day" # Since the authentication hash is valid for a short period, this is effectively per-user if settings.DEBUG: - "1000/day" + rate = "1000/day" # Throttle per wizard hash def get_cache_key(self, request, view): From 4860a561e19e0e2dd8ed2fab62002d29420b855c Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 21 Jul 2025 18:40:03 +0200 Subject: [PATCH 3/4] adjust test to not use DEBUG --- posthog/api/wizard/test/test_http.py | 1 + 1 file changed, 1 insertion(+) diff --git a/posthog/api/wizard/test/test_http.py b/posthog/api/wizard/test/test_http.py index 2c2e1bb4ade3..db7010d44a86 100644 --- a/posthog/api/wizard/test/test_http.py +++ b/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"} From fc921910b61868fa4bfa305cf19a762d4ff068ff Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 21 Jul 2025 19:28:47 +0200 Subject: [PATCH 4/4] fix failing test --- .run/PostHog.run.xml | 4 ++-- posthog/rate_limit.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/.run/PostHog.run.xml b/.run/PostHog.run.xml index 07ed9f65ef19..3e6ba9dbfffa 100644 --- a/.run/PostHog.run.xml +++ b/.run/PostHog.run.xml @@ -1,10 +1,11 @@ -