Skip to content

Commit

Permalink
tests(Analytics): update tests
Browse files Browse the repository at this point in the history
  • Loading branch information
cheikhgwane committed Jul 8, 2024
1 parent 8b2b9ec commit 608fc3d
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 136 deletions.
213 changes: 77 additions & 136 deletions hexa/core/tests/test_analytics.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import json
from datetime import datetime
from unittest import mock
from uuid import uuid4

Expand Down Expand Up @@ -49,15 +49,18 @@ def test_track_event_user_no_token(self, mock_mixpanel):
track(self.USER, event, properties, mock_request)
mock_mixpanel.assert_not_called()

def test_track_user_analytics_not_enabled(
@mock.patch("hexa.core.analytics.Mixpanel")
def test_track_pipelinerun_user_exist(
self,
mock_datetime,
mock_mixpanel,
):
self.USER.analytics_enabled = False
self.USER.save()
now = datetime.now()
mock_datetime.now.return_value = now
with self.settings(MIXPANEL_TOKEN="mixpanel_token"):
mock_mixpanel_instance = mock.MagicMock()
mock_mixpanel.return_value = mock_mixpanel_instance

mixpanel_token = "token"
# mock the flush method of the BufferedConsumer
with self.settings(MIXPANEL_TOKEN=mixpanel_token):
pipeline_run = self.PIPELINE.run(
user=self.USER,
pipeline_version=self.PIPELINE.last_version,
Expand All @@ -72,90 +75,47 @@ def test_track_user_analytics_not_enabled(
"dataset_id": str(uuid4()),
"creation_source": "SDK",
"workspace": self.WORKSPACE.slug,
"timestamp": now.timestamp(),
}

track(pipeline_run_user, event, properties)
self.assertEqual(len(mixpanel_consumer._buffers["events"]), 0)
self.assertEqual(len(mixpanel_consumer._buffers["people"]), 0)

def test_track_user_analytics_enabled(
self,
):
mixpanel_token = "token"
# mock the flush method of the BufferedConsumer
with self.settings(MIXPANEL_TOKEN=mixpanel_token), mock.patch.object(
mixpanel_consumer, "flush"
) as mocked_flush:
mock_request = mock.Mock()
mock_request.headers = {
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
}
mock_request.META = {"REMOTE_ADDR": "127.0.0.1"}

track(self.USER, "login", request=mock_request)

self.assertEqual(len(mixpanel_consumer._buffers["events"]), 1)
self.assertEqual(len(mixpanel_consumer._buffers["people"]), 1)

mixpanel_args = json.loads(mixpanel_consumer._buffers["events"][0])

self.assertTrue(mixpanel_args["event"], "login")
self.assertTrue(
mixpanel_args["properties"]["distinct_id"], str(self.USER.id)
mock_mixpanel_instance.track.assert_called_once_with(
distinct_id=str(self.USER.id), event_name=event, properties=properties
)
mock_mixpanel_instance.people_set_once.assert_called_once_with(
distinct_id=str(self.USER.id),
properties={
"$email": self.USER.email,
"$name": self.USER.display_name,
"staff_status": self.USER.is_staff,
"superuser_status": self.USER.is_superuser,
"email_domain": "bluesquarehub.com",
"features_flag": [],
},
)
self.assertTrue(mixpanel_args["properties"]["$device"], "Mac")
self.assertTrue(mixpanel_args["properties"]["$browser"], "Chrome")
self.assertTrue(mixpanel_args["properties"]["ip"], "127.0.0.1")

mocked_flush.assert_not_called()

def test_track_pipelinerun_user_exist(
@mock.patch("hexa.core.analytics.Mixpanel")
@mock.patch("hexa.core.analytics.datetime")
def test_track_user_analytics_not_enabled(
self,
mock_datetime,
mock_mixpanel,
):
mixpanel_token = "token"
# mock the flush method of the BufferedConsumer
with self.settings(MIXPANEL_TOKEN=mixpanel_token), mock.patch.object(
mixpanel_consumer, "flush"
) as mocked_flush:
pipeline_run = self.PIPELINE.run(
user=self.USER,
pipeline_version=self.PIPELINE.last_version,
trigger_mode=PipelineRunTrigger.MANUAL,
config={},
)
pipeline_run_user = PipelineRunUser(pipeline_run=pipeline_run)
now = datetime.now()
mock_datetime.now.return_value = now

event = "dataset_version_created"
properties = {
"dataset_version": "version",
"dataset_id": str(uuid4()),
"creation_source": "SDK",
"workspace": self.WORKSPACE.slug,
}

track(pipeline_run_user, event, properties)
mock_mixpanel_instance = mock.MagicMock()
mock_mixpanel.return_value = mock_mixpanel_instance

self.assertEqual(len(mixpanel_consumer._buffers["events"]), 1)
self.assertEqual(len(mixpanel_consumer._buffers["people"]), 1)
mocked_flush.assert_not_called()

mixpanel_args = json.loads(mixpanel_consumer._buffers["events"][0])
self.assertTrue(
set(properties.items()).issubset(
set(mixpanel_args["properties"].items())
)
)
self.USER.analytics_enabled = False
self.USER.save()

def test_track_pipelinerun_user_none(
self,
):
mixpanel_token = "token"
# mock the flush method of the BufferedConsumer
with self.settings(MIXPANEL_TOKEN=mixpanel_token), mock.patch.object(
mixpanel_consumer, "flush"
):
with self.settings(MIXPANEL_TOKEN=mixpanel_token):
pipeline_run = self.PIPELINE.run(
user=None,
user=self.USER,
pipeline_version=self.PIPELINE.last_version,
trigger_mode=PipelineRunTrigger.MANUAL,
config={},
Expand All @@ -171,73 +131,58 @@ def test_track_pipelinerun_user_none(
}

track(pipeline_run_user, event, properties)
self.assertEqual(len(mixpanel_consumer._buffers["events"]), 1)
self.assertEqual(len(mixpanel_consumer._buffers["people"]), 0)

mixpanel_args = json.loads(mixpanel_consumer._buffers["events"][0])

self.assertTrue(mixpanel_args["event"], "login")
self.assertEqual(mixpanel_args["properties"]["distinct_id"], None)
mock_mixpanel_instance.assert_not_called()

def test_track_pipelinerun_user_with_request(
@mock.patch("hexa.core.analytics.Mixpanel")
@mock.patch("hexa.core.analytics.datetime")
def test_track_user_analytics_enabled(
self,
mock_datetime,
mock_mixpanel,
):
mixpanel_token = "token"
# mock the flush method of the BufferedConsumer
with self.settings(MIXPANEL_TOKEN=mixpanel_token), mock.patch.object(
mixpanel_consumer, "flush"
):
pipeline_run = self.PIPELINE.run(
user=None,
pipeline_version=self.PIPELINE.last_version,
trigger_mode=PipelineRunTrigger.MANUAL,
config={},
)
pipeline_run_user = PipelineRunUser(pipeline_run=pipeline_run)

event = "dataset_version_created"
properties = {
"dataset_version": "version",
"dataset_id": str(uuid4()),
"creation_source": "SDK",
"workspace": self.WORKSPACE.slug,
}
mock_request = mock.Mock()
mock_request.headers = {
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
}
mock_request.META = {"REMOTE_ADDR": "127.0.0.1"}

track(pipeline_run_user, event, properties, mock_request)
now = datetime.now()
mock_datetime.now.return_value = now

self.assertEqual(len(mixpanel_consumer._buffers["events"]), 1)
self.assertEqual(len(mixpanel_consumer._buffers["people"]), 0)
mock_mixpanel_instance = mock.MagicMock()
mock_mixpanel.return_value = mock_mixpanel_instance
mixpanel_token = "token"

mixpanel_args = json.loads(mixpanel_consumer._buffers["events"][0])
self.USER.analytics_enabled = True
self.USER.save()

self.assertTrue(mixpanel_args["event"], "login")
self.assertEqual(mixpanel_args["properties"]["distinct_id"], None)
self.assertFalse(
all(
[
f in mixpanel_args["properties"]
for f in ["$browser", "$device", "$os", "ip"]
]
),
with self.settings(MIXPANEL_TOKEN=mixpanel_token):
track(self.USER, "login")
mock_mixpanel_instance.track.assert_called_once_with(
distinct_id=str(self.USER.id),
event_name="login",
properties={"timestamp": now.timestamp()},
)
mock_mixpanel_instance.people_set_once.assert_called_once_with(
distinct_id=str(self.USER.id),
properties={
"$email": self.USER.email,
"$name": self.USER.display_name,
"staff_status": self.USER.is_staff,
"superuser_status": self.USER.is_superuser,
"email_domain": "bluesquarehub.com",
"features_flag": [],
},
)

def test_track_pipelinerun_user_flushed_buffer(
@mock.patch("hexa.core.analytics.Mixpanel")
def test_track_pipelinerun_user_none(
self,
mock_mixpanel,
):
mock_mixpanel_instance = mock.MagicMock()
mock_mixpanel.return_value = mock_mixpanel_instance
mixpanel_token = "token"
# mock the flush method of the BufferedConsumer
with self.settings(MIXPANEL_TOKEN=mixpanel_token), mock.patch.object(
mixpanel_consumer, "_flush_endpoint"
) as mocked_flush_endpoint, mock.patch.object(
mixpanel_consumer, "_max_size", 2
mixpanel_consumer, "flush"
):
pipeline_run = self.PIPELINE.run(
user=self.USER,
user=None,
pipeline_version=self.PIPELINE.last_version,
trigger_mode=PipelineRunTrigger.MANUAL,
config={},
Expand All @@ -253,11 +198,7 @@ def test_track_pipelinerun_user_flushed_buffer(
}

track(pipeline_run_user, event, properties)
self.assertEqual(len(mixpanel_consumer._buffers["events"]), 1)
self.assertEqual(len(mixpanel_consumer._buffers["people"]), 1)

mocked_flush_endpoint.assert_not_called()

# second call to flush the buffer endpoints (events,people)
track(pipeline_run_user, event, properties)
mocked_flush_endpoint.assert_called()
mock_mixpanel_instance.track.assert_called_once_with(
distinct_id=None, event_name=event, properties=properties
)
1 change: 1 addition & 0 deletions hexa/datasets/schema/queries.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ def resolve_dataset(_, info, **kwargs):
request = info.context["request"]
try:
ds = Dataset.objects.filter_for_user(request.user).get(id=kwargs["id"])
print(f"DS {ds}", flush=True)
# triggered multiple times
# if ds.latest_version:
# properties = {
Expand Down

0 comments on commit 608fc3d

Please sign in to comment.