Skip to content

Commit

Permalink
fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@raquelmsmith
raquelmsmith committed Sep 14, 2023
1 parent 9c7ccef commit fd25902
Showing 1 changed file with 67 additions and 90 deletions.
157 changes: 67 additions & 90 deletions posthog/tasks/test/test_usage_report.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,10 @@
from posthog.session_recordings.test.test_factory import create_snapshot
from posthog.tasks.usage_report import (
_get_all_org_reports,
_get_all_usage_data_as_team_rows,
_get_full_org_usage_report,
_get_full_org_usage_report_as_dict,
_get_team_report,
capture_event,
get_instance_metadata,
send_all_org_usage_reports,
Expand Down Expand Up @@ -303,16 +305,20 @@ def _test_usage_report(self) -> List[dict]:
self._create_plugin("Installed but not enabled", False)
self._create_plugin("Installed and enabled", True)

all_reports = send_all_org_usage_reports(dry_run=False)
period = get_previous_day()
period_start, period_end = period
all_reports = _get_all_org_reports(period_start, period_end)
report = _get_full_org_usage_report_as_dict(
_get_full_org_usage_report(all_reports[str(self.organization.id)], get_instance_metadata(period))
)

report = all_reports[0]
assert report["table_sizes"]
assert report["table_sizes"]["posthog_event"] < 10**7 # <10MB
assert report["table_sizes"]["posthog_sessionrecordingevent"] < 10**7 # <10MB

assert len(all_reports) == 2

expectation = [
expectations = [
{
"deployment_infrastructure": "tests",
"realm": "hosted-clickhouse",
Expand All @@ -323,12 +329,12 @@ def _test_usage_report(self) -> List[dict]:
"site_url": "http://test.posthog.com",
"product": "open source",
"helm": {},
"clickhouse_version": all_reports[0]["clickhouse_version"],
"clickhouse_version": report["clickhouse_version"],
"users_who_logged_in": [],
"users_who_logged_in_count": 0,
"users_who_signed_up": [],
"users_who_signed_up_count": 0,
"table_sizes": all_reports[0]["table_sizes"],
"table_sizes": report["table_sizes"],
"plugins_installed": {"Installed and enabled": 1, "Installed but not enabled": 1},
"plugins_enabled": {"Installed and enabled": 1},
"instance_tag": "none",
Expand Down Expand Up @@ -448,12 +454,12 @@ def _test_usage_report(self) -> List[dict]:
"site_url": "http://test.posthog.com",
"product": "open source",
"helm": {},
"clickhouse_version": all_reports[1]["clickhouse_version"],
"clickhouse_version": report["clickhouse_version"],
"users_who_logged_in": [],
"users_who_logged_in_count": 0,
"users_who_signed_up": [],
"users_who_signed_up_count": 0,
"table_sizes": all_reports[1]["table_sizes"],
"table_sizes": report["table_sizes"],
"plugins_installed": {"Installed and enabled": 1, "Installed but not enabled": 1},
"plugins_enabled": {"Installed and enabled": 1},
"instance_tag": "none",
Expand Down Expand Up @@ -532,18 +538,22 @@ def _test_usage_report(self) -> List[dict]:
},
]

for item in expectation:
for item in expectations:
item.update(**self.expected_properties)

# tricky: list could be in different order
assert len(all_reports) == 2
for report in all_reports:
if report["organization_id"] == expectation[0]["organization_id"]:
assert report == expectation[0]
elif report["organization_id"] == expectation[1]["organization_id"]:
assert report == expectation[1]
full_reports = []
for expectation in expectations:
report = _get_full_org_usage_report_as_dict(
_get_full_org_usage_report(
all_reports[expectation["organization_id"]], get_instance_metadata(period)
)
)
assert report == expectation
full_reports.append(report)

return all_reports
return full_reports

@freeze_time("2022-01-10T00:01:00Z")
@patch("os.environ", {"DEPLOYMENT": "tests"})
Expand All @@ -559,6 +569,8 @@ def test_unlicensed_usage_report(self, mock_post: MagicMock, mock_client: MagicM
mock_client.return_value = mock_posthog

all_reports = self._test_usage_report()
with self.settings(SITE_URL="http://test.posthog.com"):
send_all_org_usage_reports()

# Check calls to other services
mock_post.assert_not_called()
Expand Down Expand Up @@ -604,20 +616,21 @@ def test_usage_report_hogql_queries(self) -> None:
run_events_query(query=EventsQuery(select=["event"], limit=50), team=self.team)
sync_execute("SYSTEM FLUSH LOGS")

all_reports = send_all_org_usage_reports(dry_run=False, at=str(now() + relativedelta(days=1)))
assert len(all_reports) == 1
period = get_previous_day(at=now() + relativedelta(days=1))
period_start, period_end = period
all_reports = _get_all_usage_data_as_team_rows(period_start, period_end)

report = all_reports[0]["teams"][str(self.team.pk)]
report = _get_team_report(all_reports, self.team)

# We selected 200 or 50 rows, but still read 100 rows to return the query
assert report["hogql_app_rows_read"] == 100
assert report["hogql_app_bytes_read"] > 0
assert report["event_explorer_app_rows_read"] == 100
assert report["event_explorer_app_bytes_read"] > 0
assert report.hogql_app_rows_read == 100
assert report.hogql_app_bytes_read > 0
assert report.event_explorer_app_rows_read == 100
assert report.event_explorer_app_bytes_read > 0

# Nothing was read via the API
assert report["hogql_api_rows_read"] == 0
assert report["event_explorer_api_rows_read"] == 0
assert report.hogql_api_rows_read == 0
assert report.event_explorer_api_rows_read == 0


@freeze_time("2022-01-10T00:01:00Z")
Expand Down Expand Up @@ -687,21 +700,19 @@ def test_usage_report_decide_requests(self, billing_task_mock: MagicMock, postho
flush_persons_and_events()

with self.settings(DECIDE_BILLING_ANALYTICS_TOKEN="correct"):
all_reports = send_all_org_usage_reports(dry_run=False, at=str(now() + relativedelta(days=1)))
period = get_previous_day(at=now() + relativedelta(days=1))
period_start, period_end = period
all_reports = _get_all_org_reports(period_start, period_end)

assert len(all_reports) == 3

all_reports = sorted(all_reports, key=lambda x: x["organization_name"])

assert [all_reports["organization_name"] for all_reports in all_reports] == [
"Org 1",
"Org 2",
"PostHog",
]

org_1_report = all_reports[0]
org_2_report = all_reports[1]
analytics_report = all_reports[2]
org_1_report = _get_full_org_usage_report_as_dict(
_get_full_org_usage_report(all_reports[str(self.org_1.id)], get_instance_metadata(period))
)
assert org_1_report["organization_name"] == "Org 1"
org_2_report = _get_full_org_usage_report_as_dict(
_get_full_org_usage_report(all_reports[str(self.org_2.id)], get_instance_metadata(period))
)

assert org_1_report["organization_name"] == "Org 1"
assert org_1_report["decide_requests_count_in_period"] == 11
Expand All @@ -728,26 +739,6 @@ def test_usage_report_decide_requests(self, billing_task_mock: MagicMock, postho
assert org_2_report["teams"]["5"]["billable_feature_flag_requests_count_in_period"] == 0
assert org_2_report["teams"]["5"]["billable_feature_flag_requests_count_in_month"] == 0

# billing service calls are made only for org1, which has decide requests, and analytics org - which has decide usage events.
calls = [
call(
org_1_report["organization_id"],
ANY,
),
call(
analytics_report["organization_id"],
ANY,
),
]
assert billing_task_mock.delay.call_count == 2
billing_task_mock.delay.assert_has_calls(
calls,
any_order=True,
)

# capture usage report calls are made for all orgs
assert posthog_capture_mock.return_value.capture.call_count == 3

@patch("posthog.tasks.usage_report.Client")
@patch("posthog.tasks.usage_report.send_report_to_billing_service")
def test_usage_report_local_evaluation_requests(
Expand Down Expand Up @@ -799,21 +790,19 @@ def test_usage_report_local_evaluation_requests(
flush_persons_and_events()

with self.settings(DECIDE_BILLING_ANALYTICS_TOKEN="correct"):
all_reports = send_all_org_usage_reports(dry_run=False, at=str(now() + relativedelta(days=1)))
period = get_previous_day(at=now() + relativedelta(days=1))
period_start, period_end = period
all_reports = _get_all_org_reports(period_start, period_end)

assert len(all_reports) == 3

all_reports = sorted(all_reports, key=lambda x: x["organization_name"])

assert [all_reports["organization_name"] for all_reports in all_reports] == [
"Org 1",
"Org 2",
"PostHog",
]

org_1_report = all_reports[0]
org_2_report = all_reports[1]
analytics_report = all_reports[2]
org_1_report = _get_full_org_usage_report_as_dict(
_get_full_org_usage_report(all_reports[str(self.org_1.id)], get_instance_metadata(period))
)
assert org_1_report["organization_name"] == "Org 1"
org_2_report = _get_full_org_usage_report_as_dict(
_get_full_org_usage_report(all_reports[str(self.org_2.id)], get_instance_metadata(period))
)

assert org_1_report["organization_name"] == "Org 1"
assert org_1_report["local_evaluation_requests_count_in_period"] == 11
Expand Down Expand Up @@ -844,26 +833,6 @@ def test_usage_report_local_evaluation_requests(
assert org_2_report["teams"]["5"]["billable_feature_flag_requests_count_in_period"] == 0
assert org_2_report["teams"]["5"]["billable_feature_flag_requests_count_in_month"] == 0

# billing service calls are made only for org1, which has decide requests, and analytics org - which has local evaluation usage events.
calls = [
call(
org_1_report["organization_id"],
ANY,
),
call(
analytics_report["organization_id"],
ANY,
),
]
assert billing_task_mock.delay.call_count == 2
billing_task_mock.delay.assert_has_calls(
calls,
any_order=True,
)

# capture usage report calls are made for all orgs
assert posthog_capture_mock.return_value.capture.call_count == 3


class SendUsageTest(LicensedTestMixin, ClickhouseDestroyTablesMixin, APIBaseTest):
def setUp(self) -> None:
Expand Down Expand Up @@ -950,18 +919,26 @@ def test_send_usage_cloud(self, mock_post: MagicMock, mock_client: MagicMock) ->
mock_posthog = MagicMock()
mock_client.return_value = mock_posthog

all_reports = send_all_org_usage_reports(dry_run=False)
period = get_previous_day()
period_start, period_end = period
all_reports = _get_all_org_reports(period_start, period_end)
full_report_as_dict = _get_full_org_usage_report_as_dict(
_get_full_org_usage_report(all_reports[str(self.organization.id)], get_instance_metadata(period))
)
send_all_org_usage_reports(dry_run=False)
license = License.objects.first()
assert license
token = build_billing_token(license, self.organization)
mock_post.assert_called_once_with(
f"{BILLING_SERVICE_URL}/api/usage", json=all_reports[0], headers={"Authorization": f"Bearer {token}"}
f"{BILLING_SERVICE_URL}/api/usage",
json=full_report_as_dict,
headers={"Authorization": f"Bearer {token}"},
)

mock_posthog.capture.assert_any_call(
self.user.distinct_id,
"organization usage report",
{**all_reports[0], "scope": "user"},
{**full_report_as_dict, "scope": "user"},
groups={"instance": "http://localhost:8000", "organization": str(self.organization.id)},
timestamp=None,
)
Expand Down

0 comments on commit fd25902

Please sign in to comment.