From 4bb6f727c06baf9199b1b7bb8bf648a912d34cf5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Mon, 18 Sep 2023 12:02:40 +0200 Subject: [PATCH] fix: Fix tests to assert for isoformat --- posthog/temporal/tests/batch_exports/base.py | 8 +++++++ .../tests/batch_exports/test_batch_exports.py | 24 ++++++++++++++++--- .../test_s3_batch_export_workflow.py | 20 +++++++++++----- .../test_snowflake_batch_export_workflow.py | 3 ++- 4 files changed, 45 insertions(+), 10 deletions(-) diff --git a/posthog/temporal/tests/batch_exports/base.py b/posthog/temporal/tests/batch_exports/base.py index b1ab1caf909099..88a52fe798426e 100644 --- a/posthog/temporal/tests/batch_exports/base.py +++ b/posthog/temporal/tests/batch_exports/base.py @@ -1,3 +1,4 @@ +import datetime as dt import json import typing @@ -71,3 +72,10 @@ async def insert_events(client: ClickHouseClient, events: list[EventValues]): def amaterialize(table: typing.Literal["events", "person", "groups"], column: str): """Materialize a column in a table.""" return materialize(table, column) + + +def to_isoformat(d: str | None) -> str | None: + """Parse a string and return it as default isoformatted.""" + if d is None: + return None + return dt.datetime.fromisoformat(d).replace(tzinfo=dt.timezone.utc).isoformat() diff --git a/posthog/temporal/tests/batch_exports/test_batch_exports.py b/posthog/temporal/tests/batch_exports/test_batch_exports.py index 913cd5b45d2c36..50ee763b5d4d9b 100644 --- a/posthog/temporal/tests/batch_exports/test_batch_exports.py +++ b/posthog/temporal/tests/batch_exports/test_batch_exports.py @@ -12,6 +12,9 @@ import pytest_asyncio from django.conf import settings +from posthog.temporal.tests.batch_exports.base import ( + to_isoformat, +) from posthog.temporal.workflows.batch_exports import ( BatchExportTemporaryFile, get_data_interval, @@ -288,8 +291,13 @@ async def test_get_results_iterator(client): for expected, result in zip(all_expected, all_result): for key, value in result.items(): + if key in ("timestamp", "inserted_at", "created_at"): + expected_value = to_isoformat(expected[key]) + else: + expected_value = expected[key] + # Some keys will be missing from result, so let's only check the ones we have. - assert value == expected[key], f"{key} value in {result} didn't match value in {expected}" + assert value == expected_value, f"{key} value in {result} didn't match value in {expected}" @pytest.mark.django_db @@ -343,8 +351,13 @@ async def test_get_results_iterator_handles_duplicates(client): for expected, result in zip(all_expected, all_result): for key, value in result.items(): + if key in ("timestamp", "inserted_at", "created_at"): + expected_value = to_isoformat(expected[key]) + else: + expected_value = expected[key] + # Some keys will be missing from result, so let's only check the ones we have. - assert value == expected[key], f"{key} value in {result} didn't match value in {expected}" + assert value == expected_value, f"{key} value in {result} didn't match value in {expected}" @pytest.mark.django_db @@ -400,8 +413,13 @@ async def test_get_results_iterator_can_exclude_events(client): for expected, result in zip(all_expected, all_result): for key, value in result.items(): + if key in ("timestamp", "inserted_at", "created_at"): + expected_value = to_isoformat(expected[key]) + else: + expected_value = expected[key] + # Some keys will be missing from result, so let's only check the ones we have. - assert value == expected[key], f"{key} value in {result} didn't match value in {expected}" + assert value == expected_value, f"{key} value in {result} didn't match value in {expected}" @pytest.mark.parametrize( diff --git a/posthog/temporal/tests/batch_exports/test_s3_batch_export_workflow.py b/posthog/temporal/tests/batch_exports/test_s3_batch_export_workflow.py index b4e51bc9f8b8eb..08f0d285a944c1 100644 --- a/posthog/temporal/tests/batch_exports/test_s3_batch_export_workflow.py +++ b/posthog/temporal/tests/batch_exports/test_s3_batch_export_workflow.py @@ -25,6 +25,7 @@ EventValues, amaterialize, insert_events, + to_isoformat, ) from posthog.temporal.tests.batch_exports.fixtures import ( acreate_batch_export, @@ -124,12 +125,19 @@ def assert_events_in_s3( if exclude_events is None: exclude_events = [] - expected_events = [ - {k: v for k, v in event.items() if k not in ["team_id", "_timestamp"]} - for event in events - if event["event"] not in exclude_events - ] - expected_events.sort(key=lambda x: x["timestamp"]) + def to_expected_event(event): + mapping_functions = { + "timestamp": to_isoformat, + "inserted_at": to_isoformat, + "created_at": to_isoformat, + } + return { + k: mapping_functions.get(k, lambda x: x)(v) for k, v in event.items() if k not in ["team_id", "_timestamp"] + } + + expected_events = list(map(to_expected_event, (event for event in events if event["event"] not in exclude_events))) + + expected_events.sort(key=lambda x: x["timestamp"] if x["timestamp"] is not None else 0) # First check one event, the first one, so that we can get a nice diff if # the included data is different. diff --git a/posthog/temporal/tests/batch_exports/test_snowflake_batch_export_workflow.py b/posthog/temporal/tests/batch_exports/test_snowflake_batch_export_workflow.py index 3b153668d55495..3e3aefdb6ad893 100644 --- a/posthog/temporal/tests/batch_exports/test_snowflake_batch_export_workflow.py +++ b/posthog/temporal/tests/batch_exports/test_snowflake_batch_export_workflow.py @@ -21,6 +21,7 @@ from posthog.temporal.tests.batch_exports.base import ( EventValues, insert_events, + to_isoformat, ) from posthog.temporal.tests.batch_exports.fixtures import ( acreate_batch_export, @@ -361,7 +362,7 @@ async def test_snowflake_export_workflow_exports_events_in_the_last_hour_for_the { "uuid": event["uuid"], "event": event["event"], - "timestamp": event["timestamp"], + "timestamp": to_isoformat(event["timestamp"]), "properties": event["properties"], "person_id": event["person_id"], }