fix typing

PostHog · Oct 17, 2024 · f0bc34a · f0bc34a
1 parent 67c741d
commit f0bc34a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/posthog/hogql_queries/experiments/experiment_funnel_query_runner.py b/posthog/hogql_queries/experiments/experiment_funnel_query_runner.py
@@ -152,12 +152,13 @@ def _validate_event_variants(self, funnel_results: FunnelsQueryResponse):
         eventsWithOrderZero = []
         for eventArr in funnel_results.results:
             for event in eventArr:
-                if event.get("order") == 0:
-                    eventsWithOrderZero.append(event)
+                event_dict = cast(dict[str, Any], event)
+                if event_dict.get("order") == 0:
+                    eventsWithOrderZero.append(event_dict)
 
         # Check if "control" is present
         for event in eventsWithOrderZero:
-            event_variant = event.get("breakdown_value")[0]
+            event_variant = event.get("breakdown_value", [None])[0]
             if event_variant == "control":
                 errors[ExperimentNoResultsErrorKeys.NO_CONTROL_VARIANT] = False
                 errors[ExperimentNoResultsErrorKeys.NO_FLAG_INFO] = False
@@ -166,7 +167,7 @@ def _validate_event_variants(self, funnel_results: FunnelsQueryResponse):
         # Check if at least one of the test variants is present
         test_variants = [variant for variant in self.variants if variant != "control"]
         for event in eventsWithOrderZero:
-            event_variant = event.get("breakdown_value")[0]
+            event_variant = event.get("breakdown_value", [None])[0]
             if event_variant in test_variants:
                 errors[ExperimentNoResultsErrorKeys.NO_TEST_VARIANT] = False
                 errors[ExperimentNoResultsErrorKeys.NO_FLAG_INFO] = False

diff --git a/posthog/hogql_queries/experiments/test/test_experiment_funnel_query_runner.py b/posthog/hogql_queries/experiments/test/test_experiment_funnel_query_runner.py
@@ -1,3 +1,4 @@
+from typing import cast
 from posthog.hogql_queries.experiments.experiment_funnel_query_runner import ExperimentFunnelQueryRunner
 from posthog.models.experiment import Experiment
 from posthog.models.feature_flag.feature_flag import FeatureFlag
@@ -227,7 +228,7 @@ def test_validate_event_variants_no_events(self):
                 ExperimentNoResultsErrorKeys.NO_TEST_VARIANT: True,
             }
         )
-        self.assertEqual(context.exception.detail[0], expected_errors)
+        self.assertEqual(cast(list, context.exception.detail)[0], expected_errors)
 
     @freeze_time("2020-01-01T12:00:00Z")
     def test_validate_event_variants_no_control(self):
@@ -269,7 +270,7 @@ def test_validate_event_variants_no_control(self):
                 ExperimentNoResultsErrorKeys.NO_TEST_VARIANT: False,
             }
         )
-        self.assertEqual(context.exception.detail[0], expected_errors)
+        self.assertEqual(cast(list, context.exception.detail)[0], expected_errors)
 
     @freeze_time("2020-01-01T12:00:00Z")
     def test_validate_event_variants_no_test(self):
@@ -311,7 +312,7 @@ def test_validate_event_variants_no_test(self):
                 ExperimentNoResultsErrorKeys.NO_TEST_VARIANT: True,
             }
         )
-        self.assertEqual(context.exception.detail[0], expected_errors)
+        self.assertEqual(cast(list, context.exception.detail)[0], expected_errors)
 
     @freeze_time("2020-01-01T12:00:00Z")
     def test_validate_event_variants_no_flag_info(self):
@@ -355,4 +356,4 @@ def test_validate_event_variants_no_flag_info(self):
                 ExperimentNoResultsErrorKeys.NO_TEST_VARIANT: True,
             }
         )
-        self.assertEqual(context.exception.detail[0], expected_errors)
+        self.assertEqual(cast(list, context.exception.detail)[0], expected_errors)