test(hogql): add tests for hogql query cache

PostHog · Sep 19, 2023 · fdcd908 · fdcd908
1 parent c61a9d0
commit fdcd908
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 3 deletions.
diff --git a/posthog/hogql_queries/query_runner.py b/posthog/hogql_queries/query_runner.py
@@ -74,8 +74,7 @@ def run(self, refresh_requested: bool) -> CachedQueryResponse:
 
         if not refresh_requested:
             cached_response = get_safe_cache(cache_key)
-
-            if cached_response and cached_response.result:
+            if cached_response:
                 if not self._is_stale(cached_response):
                     QUERY_CACHE_HIT_COUNTER.labels(team_id=self.team.pk, cache_hit="hit").inc()
                     cached_response.is_cached = True

diff --git a/posthog/hogql_queries/test/test_query_runner.py b/posthog/hogql_queries/test/test_query_runner.py
@@ -1,6 +1,10 @@
+from datetime import datetime, timedelta
+from dateutil.parser import isoparse
+from zoneinfo import ZoneInfo
 from typing import Any, List, Literal, Optional, Type
+from freezegun import freeze_time
 from pydantic import BaseModel
-from posthog.hogql_queries.query_runner import QueryRunner
+from posthog.hogql_queries.query_runner import QueryResponse, QueryRunner
 from posthog.models.team.team import Team
 from posthog.test.base import BaseTest
 from posthog.types import InsightQueryNode
@@ -19,6 +23,17 @@ def setup_test_query_runner_class(self, query_class: Type[InsightQueryNode] = Te
         class TestQueryRunner(QueryRunner):
             query_type = query_class
 
+            def calculate(self) -> QueryResponse:
+                return QueryResponse(result=list())
+
+            def _refresh_frequency(self) -> timedelta:
+                return timedelta(minutes=4)
+
+            def _is_stale(self, cached_result_package) -> bool:
+                return isoparse(cached_result_package.last_refresh) + timedelta(minutes=10) <= datetime.now(
+                    tz=ZoneInfo("UTC")
+                )
+
         TestQueryRunner.__abstractmethods__ = frozenset()
 
         return TestQueryRunner
@@ -83,3 +98,28 @@ def test_cache_key_different_timezone(self):
 
         cache_key = runner._cache_key()
         self.assertEqual(cache_key, "cache_0fa2172980705adb41741351f40189b7")
+
+    def test_cache_response(self):
+        TestQueryRunner = self.setup_test_query_runner_class()
+
+        runner = TestQueryRunner(query={"some_attr": "bla"}, team=self.team)  # type: ignore
+
+        with freeze_time(datetime(2023, 2, 4, 13, 37, 42)):
+            # returns fresh response if uncached
+            response = runner.run(refresh_requested=False)
+            self.assertEqual(response.is_cached, False)
+            self.assertEqual(response.last_refresh, "2023-02-04T13:37:42Z")
+            self.assertEqual(response.next_allowed_client_refresh, "2023-02-04T13:41:42Z")
+
+            # returns cached response afterwards
+            response = runner.run(refresh_requested=False)
+            self.assertEqual(response.is_cached, True)
+
+            # return fresh response if refresh requested
+            response = runner.run(refresh_requested=True)
+            self.assertEqual(response.is_cached, False)
+
+        with freeze_time(datetime(2023, 2, 4, 13, 37 + 11, 42)):
+            # returns fresh response if stale
+            response = runner.run(refresh_requested=False)
+            self.assertEqual(response.is_cached, False)