chore: Move worker billing functionality to shared (#436)

* chore: Move worker billing functionality to shared

* why the log?

* update the tests

* update the tests 2

* Refactor tests

* Update to fn based tests, easier and simpler to read

shared/billing/__init__.py

@@ -1,15 +1,23 @@
 from enum import Enum
 
+from django.conf import settings
+
+from shared.license import get_current_license
+
 
 class BillingPlan(Enum):
+    users_ghm = "users"
     users_monthly = "users-inappm"
     users_yearly = "users-inappy"
     users_free = "users-free"
     users_basic = "users-basic"
+    users_trial = "users-trial"
     pr_monthly = "users-pr-inappm"
     pr_yearly = "users-pr-inappy"
     enterprise_cloud_yearly = "users-enterprisey"
     enterprise_cloud_monthly = "users-enterprisem"
+    team_monthly = "users-teamm"
+    team_yearly = "users-teamy"
 
     def __init__(self, db_name):
         self.db_name = db_name
@@ -26,3 +34,21 @@ def is_enterprise_cloud_plan(plan: BillingPlan) -> bool:
         BillingPlan.enterprise_cloud_monthly,
         BillingPlan.enterprise_cloud_yearly,
     ]
+
+
+def is_pr_billing_plan(plan: str) -> bool:
+    if not settings.IS_ENTERPRISE:
+        return plan in [
+            BillingPlan.pr_monthly.value,
+            BillingPlan.pr_yearly.value,
+            BillingPlan.users_free.value,
+            BillingPlan.users_basic.value,
+            BillingPlan.users_trial.value,
+            BillingPlan.enterprise_cloud_monthly.value,
+            BillingPlan.enterprise_cloud_yearly.value,
+            BillingPlan.team_monthly.value,
+            BillingPlan.team_yearly.value,
+            BillingPlan.users_ghm.value,
+        ]
+    else:
+        return get_current_license().is_pr_billing

tests/unit/billing/test_enum_definitions.py

@@ -1,4 +1,64 @@
-from shared.billing import BillingPlan, is_enterprise_cloud_plan
+import pytest
+from django.test import override_settings
+
+from shared.billing import BillingPlan, is_enterprise_cloud_plan, is_pr_billing_plan
+from shared.django_apps.codecov_auth.tests.factories import OwnerFactory
+
+
+@pytest.fixture
+def dbsession(db):
+    return db
+
+
+@override_settings(IS_ENTERPRISE=False)
+def test_pr_author_plan_check(dbsession, mock_configuration):
+    owner = OwnerFactory(service="github", plan="users-pr-inappm")
+    if dbsession is not None:
+        dbsession.add(owner)
+        dbsession.flush()
+    assert is_pr_billing_plan(owner.plan)
+
+
+@override_settings(IS_ENTERPRISE=True)
+def test_pr_author_enterprise_plan_check(dbsession, mock_configuration):
+    owner = OwnerFactory(service="github")
+    if dbsession is not None:
+        dbsession.add(owner)
+        dbsession.flush()
+
+    encrypted_license = "wxWEJyYgIcFpi6nBSyKQZQeaQ9Eqpo3SXyUomAqQOzOFjdYB3A8fFM1rm+kOt2ehy9w95AzrQqrqfxi9HJIb2zLOMOB9tSy52OykVCzFtKPBNsXU/y5pQKOfV7iI3w9CHFh3tDwSwgjg8UsMXwQPOhrpvl2GdHpwEhFdaM2O3vY7iElFgZfk5D9E7qEnp+WysQwHKxDeKLI7jWCnBCBJLDjBJRSz0H7AfU55RQDqtTrnR+rsLDHOzJ80/VxwVYhb"
+    mock_configuration.params["setup"]["enterprise_license"] = encrypted_license
+    mock_configuration.params["setup"]["codecov_dashboard_url"] = (
+        "https://codecov.mysite.com"
+    )
+
+    assert is_pr_billing_plan(owner.plan)
+
+
+@override_settings(IS_ENTERPRISE=False)
+def test_plan_not_pr_author(dbsession, mock_configuration):
+    owner = OwnerFactory(service="github", plan=BillingPlan.users_monthly.value)
+    if dbsession is not None:
+        dbsession.add(owner)
+        dbsession.flush()
+
+    assert not is_pr_billing_plan(owner.plan)
+
+
+@override_settings(IS_ENTERPRISE=True)
+def test_pr_author_enterprise_plan_check_non_pr_plan(dbsession, mock_configuration):
+    owner = OwnerFactory(service="github")
+    if dbsession is not None:
+        dbsession.add(owner)
+        dbsession.flush()
+
+    encrypted_license = "0dRbhbzp8TVFQp7P4e2ES9lSfyQlTo8J7LQ"
+    mock_configuration.params["setup"]["enterprise_license"] = encrypted_license
+    mock_configuration.params["setup"]["codecov_dashboard_url"] = (
+        "https://codeov.mysite.com"
+    )
+
+    assert not is_pr_billing_plan(owner.plan)
 
 
 def test_billing_enums():
@@ -10,6 +70,8 @@ def test_billing_enums():
     assert BillingPlan.pr_yearly.db_name == "users-pr-inappy"
     assert BillingPlan.enterprise_cloud_yearly.db_name == "users-enterprisey"
     assert BillingPlan.enterprise_cloud_monthly.db_name == "users-enterprisem"
+    assert BillingPlan.team_monthly.db_name == "users-teamm"
+    assert BillingPlan.team_yearly.db_name == "users-teamy"
 
 
 def test_get_from_string():
@@ -26,6 +88,8 @@ def test_get_from_string():
         BillingPlan.from_str("users-enterprisem")
         == BillingPlan.enterprise_cloud_monthly
     )
+    assert BillingPlan.from_str("users-teamm") == BillingPlan.team_monthly
+    assert BillingPlan.from_str("users-teamy") == BillingPlan.team_yearly
 
 
 def test_is_enterprise_cloud_plan():