posthog/ee/tasks/test/test_send_license_usage.py

from unittest.mock import ANY, Mock, patch

from freezegun import freeze_time

from ee.api.test.base import LicensedTestMixin
from ee.models.license import License
from ee.tasks.send_license_usage import send_license_usage
from posthog.models.team import Team
from posthog.test.base import APIBaseTest, ClickhouseDestroyTablesMixin, _create_event, flush_persons_and_events


class SendLicenseUsageTest(LicensedTestMixin, ClickhouseDestroyTablesMixin, APIBaseTest):
    @freeze_time("2021-10-10T23:01:00Z")
    @patch("posthoganalytics.capture")
    @patch("requests.post")
    def test_send_license_usage(self, mock_post, mock_capture):
        team2 = Team.objects.create(organization=self.organization)
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-08T14:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-09T12:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-09T13:01:01Z")
        _create_event(
            event="$$internal_metrics_shouldnt_be_billed",
            team=self.team,
            distinct_id=1,
            timestamp="2021-10-09T13:01:01Z",
        )
        _create_event(event="$pageview", team=team2, distinct_id=1, timestamp="2021-10-09T14:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-10T14:01:01Z")
        flush_persons_and_events()

        mockresponse = Mock()
        mock_post.return_value = mockresponse
        mockresponse.json = lambda: {"ok": True, "valid_until": "2021-11-10T23:01:00Z"}

        send_license_usage()
        mock_post.assert_called_once_with(
            "https://license.posthog.com/licenses/usage",
            data={"date": "2021-10-09", "key": self.license.key, "events_count": 3},
        )
        mock_capture.assert_called_once_with(
            self.user.distinct_id,
            "send license usage data",
            {"date": "2021-10-09", "events_count": 3, "license_keys": ["enterprise"], "organization_name": "Test"},
            groups={"instance": ANY, "organization": str(self.organization.id)},
        )
        self.assertEqual(License.objects.get().valid_until.isoformat(), "2021-11-10T23:01:00+00:00")

    @freeze_time("2021-10-10T23:01:00Z")
    @patch("posthoganalytics.capture")
    @patch("ee.tasks.send_license_usage.sync_execute", side_effect=Exception())
    def test_send_license_error(self, mock_post, mock_capture):
        team2 = Team.objects.create(organization=self.organization)
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-08T14:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-09T12:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-09T13:01:01Z")
        _create_event(
            event="$$internal_metrics_shouldnt_be_billed",
            team=self.team,
            distinct_id=1,
            timestamp="2021-10-09T13:01:01Z",
        )
        _create_event(event="$pageview", team=team2, distinct_id=1, timestamp="2021-10-09T14:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-10T14:01:01Z")
        flush_persons_and_events()
        send_license_usage()
        mock_capture.assert_called_once_with(
            self.user.distinct_id,
            "send license usage data error",
            {"error": "", "date": "2021-10-09", "organization_name": "Test"},
            groups={"instance": ANY, "organization": str(self.organization.id)},
        )

    @freeze_time("2021-10-10T23:01:00Z")
    @patch("posthoganalytics.capture")
    @patch("requests.post")
    def test_send_license_not_found(self, mock_post, mock_capture):
        team2 = Team.objects.create(organization=self.organization)
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-08T14:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-09T12:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-09T13:01:01Z")
        _create_event(
            event="$$internal_metrics_shouldnt_be_billed",
            team=self.team,
            distinct_id=1,
            timestamp="2021-10-09T13:01:01Z",
        )
        _create_event(event="$pageview", team=team2, distinct_id=1, timestamp="2021-10-09T14:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-10T14:01:01Z")
        flush_persons_and_events()
        flush_persons_and_events()

        mockresponse = Mock()
        mock_post.return_value = mockresponse
        mockresponse.status_code = 404
        mockresponse.ok = False
        mockresponse.json = lambda: {"code": "not_found"}
        mockresponse.content = ""

        send_license_usage()

        mock_capture.assert_called_once_with(
            self.user.distinct_id,
            "send license usage data error",
            {"error": "", "date": "2021-10-09", "organization_name": "Test", "status_code": 404, "events_count": 3},
            groups={"instance": ANY, "organization": str(self.organization.id)},
        )
        self.assertEqual(License.objects.get().valid_until.isoformat(), "2021-10-10T22:01:00+00:00")


class SendLicenseUsageNoLicenseTest(APIBaseTest):
    @freeze_time("2021-10-10T23:01:00Z")
    @patch("requests.post")
    def test_no_license(self, mock_post):
        # Same test, we just don't include the LicensedTestMixin so no license
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-08T14:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-09T12:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-09T13:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-09T14:01:01Z")
        _create_event(event="$pageview", team=self.team, distinct_id=1, timestamp="2021-10-10T14:01:01Z")

        flush_persons_and_events()

        send_license_usage()

        self.assertEqual(mock_post.call_count, 0)