Fix where clause extractor tests for v1

PostHog · Jun 20, 2024 · 1887637 · 1887637
1 parent 5397435
commit 1887637
Showing 1 changed file with 15 additions and 12 deletions.
diff --git a/posthog/hogql/database/schema/util/test/test_session_where_clause_extractor.py b/posthog/hogql/database/schema/util/test/test_session_where_clause_extractor.py
@@ -8,6 +8,7 @@
 from posthog.hogql.parser import parse_select, parse_expr
 from posthog.hogql.printer import prepare_ast_for_printing, print_prepared_ast
 from posthog.hogql.visitor import clone_expr
+from posthog.schema import SessionTableVersion
 from posthog.test.base import ClickhouseTestMixin, APIBaseTest
 
 
@@ -29,11 +30,12 @@ def parse(
     return parsed
 
 
-class TestSessionWhereClauseExtractor(ClickhouseTestMixin, APIBaseTest):
+class TestSessionWhereClauseExtractorV1(ClickhouseTestMixin, APIBaseTest):
     @property
     def inliner(self):
         team = self.team
         modifiers = create_default_modifiers_for_team(team)
+        modifiers.sessionTableVersion = SessionTableVersion.V1
         context = HogQLContext(
             team_id=team.pk,
             team=team,
@@ -280,6 +282,7 @@ class TestSessionsQueriesHogQLToClickhouse(ClickhouseTestMixin, APIBaseTest):
     def print_query(self, query: str) -> str:
         team = self.team
         modifiers = create_default_modifiers_for_team(team)
+        modifiers.sessionTableVersion = SessionTableVersion.V1
         context = HogQLContext(
             team_id=team.pk,
             team=team,
@@ -329,14 +332,14 @@ def test_join_with_events(self):
 FROM
     events
     JOIN (SELECT
-        raw_sessions.session_id AS session_id
+        sessions.session_id AS session_id
     FROM
-        raw_sessions
+        sessions
     WHERE
         and(equals(sessions.team_id, {self.team.id}), ifNull(greaterOrEquals(plus(toTimeZone(sessions.min_timestamp, %(hogql_val_0)s), toIntervalDay(3)), %(hogql_val_1)s), 0))
     GROUP BY
-        raw_sessions.session_id,
-        raw_sessions.session_id) AS sessions ON equals(events.`$session_id`, sessions.session_id)
+        sessions.session_id,
+        sessions.session_id) AS sessions ON equals(events.`$session_id`, sessions.session_id)
 WHERE
     and(equals(events.team_id, {self.team.id}), greater(toTimeZone(events.timestamp, %(hogql_val_2)s), %(hogql_val_3)s))
 GROUP BY
@@ -363,17 +366,17 @@ def test_union(self):
 FROM
     events
     LEFT JOIN (SELECT
-        dateDiff(%(hogql_val_0)s, min(raw_sessions.min_timestamp), max(raw_sessions.max_timestamp)) AS `$session_duration`,
-        raw_sessions.session_id_v7 AS session_id_v7
+        dateDiff(%(hogql_val_0)s, min(sessions.min_timestamp), max(sessions.max_timestamp)) AS `$session_duration`,
+        sessions.session_id AS session_id
     FROM
-        raw_sessions
+        sessions
     WHERE
-        and(equals(raw_sessions.team_id, {self.team.id}), ifNull(lessOrEquals(minus(toTimeZone(fromUnixTimestamp(intDiv(toUInt64(bitShiftRight(raw_sessions.session_id_v7, 80)), 1000)), %(hogql_val_1)s), toIntervalDay(3)), today()), 0))
+        and(equals(sessions.team_id, {self.team.id}), ifNull(lessOrEquals(minus(toTimeZone(sessions.min_timestamp, %(hogql_val_1)s), toIntervalDay(3)), today()), 0))
     GROUP BY
-        raw_sessions.session_id_v7,
-        raw_sessions.session_id_v7) AS events__session ON equals(toUInt128(accurateCastOrNull(events.`$session_id`, %(hogql_val_2)s)), events__session.session_id_v7)
+        sessions.session_id,
+        sessions.session_id) AS events__session ON equals(events.`$session_id`, events__session.session_id)
 WHERE
-    and(equals(events.team_id, 1785), less(toTimeZone(events.timestamp, %(hogql_val_3)s), today()))
+    and(equals(events.team_id, {self.team.id}), less(toTimeZone(events.timestamp, %(hogql_val_2)s), today()))
 LIMIT {MAX_SELECT_RETURNED_ROWS}"""
         assert expected == actual