diff --git a/posthog/tasks/test/test_usage_report.py b/posthog/tasks/test/test_usage_report.py index 8193585cffff2..636b3e76b93e9 100644 --- a/posthog/tasks/test/test_usage_report.py +++ b/posthog/tasks/test/test_usage_report.py @@ -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, @@ -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", @@ -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", @@ -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", @@ -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"}) @@ -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() @@ -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") @@ -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 @@ -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( @@ -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 @@ -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: @@ -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, )