From 5e2fbd986b62db44e5800ab1fef5da92115ad26d Mon Sep 17 00:00:00 2001 From: Marcus Kok <47163063+Marcusk19@users.noreply.github.com> Date: Tue, 8 Apr 2025 08:55:37 -0400 Subject: [PATCH] marketplace: free tier integration for reconciler (PROJQUAY-5698) (#3589) free sku integration for reconciliation worker --- data/billing.py | 8 ++ util/marketplace.py | 78 ++++++++++++++++- util/test/test_marketplace.py | 2 + workers/reconciliationworker.py | 100 +++++++++++++++------- workers/test/test_reconciliationworker.py | 17 +++- 5 files changed, 170 insertions(+), 35 deletions(-) diff --git a/data/billing.py b/data/billing.py index df71ea62f..e597bd582 100644 --- a/data/billing.py +++ b/data/billing.py @@ -372,6 +372,14 @@ PLANS = [ "sku_billing": True, "plans_page_hidden": True, }, + { + "title": "freetier", + "privateRepos": 0, + "stripeId": "not_a_stripe_plan", + "rh_sku": "MW04192", + "sku_billing": False, + "plans_page_hidden": True, + }, ] RH_SKUS = [plan["rh_sku"] for plan in PLANS if plan.get("rh_sku") is not None] diff --git a/util/marketplace.py b/util/marketplace.py index b8448d552..9fcf072dd 100644 --- a/util/marketplace.py +++ b/util/marketplace.py @@ -120,7 +120,7 @@ class RedHatSubscriptionApi(object): if now_ms < end_date: logger.debug("subscription found for %s", str(skuId)) valid_subscriptions.append(subscription) - return valid_subscriptions + return valid_subscriptions if len(valid_subscriptions) > 0 else None return None def extend_subscription(self, subscription_id, endDate): @@ -171,6 +171,7 @@ class RedHatSubscriptionApi(object): "webCustomerId": customerId, } logger.debug("Created entitlement") + try: r = requests.request( method="post", @@ -187,6 +188,30 @@ class RedHatSubscriptionApi(object): return r.status_code + def remove_entitlement(self, subscription_id): + """ + Removes subscription from user given subscription_id + """ + request_url = ( + f"{self.marketplace_endpoint}/subscription/v5/terminateSubscription/{subscription_id}" + ) + request_headers = {"Content-Type": "application/json"} + + logger.debug("Terminating subscription with id %s", subscription_id) + try: + r = requests.request( + method="post", + url=request_url, + cert=self.cert, + headers=request_headers, + verify=True, + timeout=REQUEST_TIMEOUT, + ) + except requests.exceptions.ReadTimeout: + logger.info("request to %s timed out", self.marketplace_endpoint) + return 408 + return r.status_code + def get_subscription_details(self, subscription_id): """ Return the sku and expiration date for a specific subscription @@ -269,7 +294,6 @@ class RedHatSubscriptionApi(object): # Mocked classes for unit tests - TEST_USER = { "account_number": 12345, "email": "subscriptions@devtable.com", @@ -358,6 +382,45 @@ FREE_USER = { "email": "free_user@test.com", "username": "free_user", } +PAID_USER = { + "account_number": 34567, + "email": "paid@test.com", + "username": "paid_user", + "subscriptions": [ + { + "id": 12345678, + "masterEndSystemName": "Quay", + "createdEndSystemName": "SUBSCRIPTION", + "createdDate": 1675957362000, + "lastUpdateEndSystemName": "SUBSCRIPTION", + "lastUpdateDate": 1675957362000, + "installBaseStartDate": 1707368400000, + "installBaseEndDate": 1707368399000, + "webCustomerId": 123456, + "subscriptionNumber": "12399889", + "quantity": 1, + "effectiveStartDate": 1707368400000, + "effectiveEndDate": 3813177600000, + } + ], + "free_sku": [ + { + "id": 56781234, + "masterEndSystemName": "Quay", + "createdEndSystemName": "SUBSCRIPTION", + "createdDate": 1675957362000, + "lastUpdateEndSystemName": "SUBSCRIPTION", + "lastUpdateDate": 1675957362000, + "installBaseStartDate": 1707368400000, + "installBaseEndDate": 1707368399000, + "webCustomerId": 123456, + "subscriptionNumber": "12399889", + "quantity": 1, + "effectiveStartDate": 1707368400000, + "effectiveEndDate": 3813177600000, + } + ], +} class FakeUserApi(RedHatUserApi): @@ -368,6 +431,8 @@ class FakeUserApi(RedHatUserApi): def lookup_customer_id(self, email): if email == TEST_USER["email"]: return [TEST_USER["account_number"]] + if email == PAID_USER["email"]: + return [PAID_USER["account_number"]] if email == FREE_USER["email"]: return [FREE_USER["account_number"]] if email == STRIPE_USER["email"]: @@ -391,11 +456,20 @@ class FakeSubscriptionApi(RedHatSubscriptionApi): return [TEST_USER["private_subscription"]] elif customer_id == TEST_USER["account_number"] and sku_id == "MW00584MO": return [TEST_USER["reconciled_subscription"]] + elif customer_id == PAID_USER["account_number"] and sku_id == "MW02701": + return PAID_USER["subscriptions"] + elif customer_id == PAID_USER["account_number"] and sku_id == "MW04192": + return PAID_USER["free_sku"] + elif customer_id == FREE_USER["account_number"]: + return [] return None def create_entitlement(self, customer_id, sku_id): self.subscription_created = True + def remove_entitlement(self, subscription_id): + pass + def extend_subscription(self, subscription_id, end_date): self.subscription_extended = True diff --git a/util/test/test_marketplace.py b/util/test/test_marketplace.py index 028fbe840..a060ff1ed 100644 --- a/util/test/test_marketplace.py +++ b/util/test/test_marketplace.py @@ -224,6 +224,8 @@ class TestMarketplace(unittest.TestCase): assert extended_subscription is None create_subscription_response = subscription_api.create_entitlement(12345, "sku") assert create_subscription_response == 408 + remove_entitlement_response = subscription_api.remove_entitlement(12345) + assert remove_entitlement_response == 408 @patch("requests.request") def test_user_lookup(self, requests_mock): diff --git a/workers/reconciliationworker.py b/workers/reconciliationworker.py index fa53ff65c..7b71bcb0b 100644 --- a/workers/reconciliationworker.py +++ b/workers/reconciliationworker.py @@ -7,7 +7,7 @@ from app import app from app import billing as stripe from app import marketplace_subscriptions, marketplace_users from data import model -from data.billing import RECONCILER_SKUS, get_plan +from data.billing import RECONCILER_SKUS, RH_SKUS, get_plan from data.model import entitlements from util.locking import GlobalLock, LockNotAcquiredException from workers.gunicorn_worker import GunicornWorker @@ -24,6 +24,8 @@ MILLISECONDS_IN_SECONDS = 1000 SECONDS_IN_DAYS = 86400 ONE_MONTH = 30 * SECONDS_IN_DAYS * MILLISECONDS_IN_SECONDS +FREE_TIER_SKU = "MW04192" + class ReconciliationWorker(Worker): def __init__(self): @@ -42,9 +44,9 @@ class ReconciliationWorker(Worker): users = model.user.get_active_users(include_orgs=True) - stripe_users = [user for user in users if user.stripe_id is not None] + # stripe_users = [user for user in users if user.stripe_id is not None] - for user in stripe_users: + for user in users: email = user.email model_customer_ids = entitlements.get_web_customer_ids(user.id) @@ -77,39 +79,77 @@ class ReconciliationWorker(Worker): for customer_id in model_customer_ids: if customer_id not in customer_ids: entitlements.remove_web_customer_id(user, customer_id) + # check for any subscription reconciliations + stripe_customer = None + if user.stripe_id is not None: + try: + stripe_customer = stripe.Customer.retrieve(user.stripe_id) + except stripe.error.APIConnectionError: + logger.error("Cannot connect to Stripe") + continue + except stripe.error.InvalidRequestException: + logger.warn("Invalid request for stripe_id %s", user.stripe_id) + continue - # check if we need to create a subscription for customer in RH marketplace - try: - stripe_customer = stripe.Customer.retrieve(user.stripe_id) - except stripe.error.APIConnectionError: - logger.error("Cannot connect to Stripe") - continue - except stripe.error.InvalidRequestError: - logger.warn("Invalid request for stripe_id %s", user.stripe_id) - continue - try: - subscription = stripe_customer.subscription - except AttributeError: - subscription = None - for sku_id in RECONCILER_SKUS: - if subscription is not None: - plan = get_plan(stripe_customer.subscription.plan.id) - if plan is None: - continue - if plan.get("rh_sku") == sku_id: - for customer_id in customer_ids: - subscription = marketplace_api.lookup_subscription(customer_id, sku_id) - if subscription is None: - logger.debug("Found %s to create for %s", sku_id, user.username) - marketplace_api.create_entitlement(customer_id, sku_id) - break - else: - logger.debug("User %s does not have a stripe subscription", user.username) + self._iterate_over_ids(stripe_customer, customer_ids, marketplace_api, user.username) logger.debug("Finished work for user %s", user.username) logger.info("Reconciliation worker is done") + def _iterate_over_ids(self, stripe_customer, customer_ids, marketplace_api, user=None): + """ + Iterate over each customer's web id(s) and perform appropriate reconciliation actions + """ + try: + subscription = stripe_customer.subscription + except AttributeError: + subscription = None + + for customer_id in customer_ids: + paying = False + customer_skus = self._prefetch_user_entitlements(customer_id, marketplace_api) + + if stripe_customer is not None and subscription is not None: + plan = get_plan(stripe_customer.subscription.plan.id) + if plan is not None: + # check for missing sku + paying = True + plan_sku = plan.get("rh_sku") + if plan_sku not in customer_skus: + logger.debug("Found %s to create for %s", plan_sku, user) + marketplace_api.create_entitlement(customer_id, plan_sku) + # check for free tier sku + else: + # not a stripe customer but we want to check for paying subscriptions for the + # next step + if len(customer_skus) == 1 and customer_skus[0] == FREE_TIER_SKU: + # edge case where there is only one free sku present + paying = False + else: + paying = len(customer_skus) > 0 + + # check for free-tier reconciliations + if not paying and FREE_TIER_SKU not in customer_skus: + marketplace_api.create_entitlement(customer_id, FREE_TIER_SKU) + elif paying and FREE_TIER_SKU in customer_skus: + free_tier_subscriptions = marketplace_api.lookup_subscription( + customer_id, FREE_TIER_SKU + ) + # api returns a list of subscriptions so we want to make sure we remove + # all if there's more than one + for sub in free_tier_subscriptions: + id = sub.get("id") + marketplace_api.remove_entitlement(id) + + def _prefetch_user_entitlements(self, customer_id, marketplace_api): + found_skus = [] + for sku in RH_SKUS: + subscription = marketplace_api.lookup_subscription(customer_id, sku) + if subscription is not None and len(subscription) > 0: + found_skus.append(sku) + return found_skus + def _reconcile_entitlements(self, skip_lock_for_testing=False): """ Performs reconciliation for user entitlements diff --git a/workers/test/test_reconciliationworker.py b/workers/test/test_reconciliationworker.py index 4e1096595..e0873d20b 100644 --- a/workers/test/test_reconciliationworker.py +++ b/workers/test/test_reconciliationworker.py @@ -19,7 +19,18 @@ def test_skip_free_user(initialized_db): with patch.object(marketplace_subscriptions, "create_entitlement") as mock: worker._perform_reconciliation(marketplace_users, marketplace_subscriptions) - mock.assert_not_called() + # adding the free tier + mock.assert_called_with(23456, "MW04192") + + +def test_remove_free_tier(initialized_db): + # if a user has a sku and also has a free tier, the free tier should be removed + paid_user = model.user.create_user("paid_user", "password", "paid@test.com") + paid_user.save() + marketplace_subscriptions.create_entitlement(12345, "MW04192") + with patch.object(marketplace_subscriptions, "remove_entitlement") as mock: + worker._perform_reconciliation(marketplace_users, marketplace_subscriptions) + mock.assert_called_with(56781234) # fake "free" tier subscription id mocked in marketplace.py def test_reconcile_org_user(initialized_db): @@ -77,12 +88,12 @@ def test_reconcile_different_ids(initialized_db): test_user = model.user.create_user("stripe_user", "password", "stripe_user@test.com") test_user.stripe_id = "cus_" + "".join(random.choices(string.ascii_lowercase, k=14)) test_user.save() - model.entitlements.save_web_customer_id(test_user, 12345) + model.entitlements.save_web_customer_id(test_user, 55555) worker._perform_reconciliation(marketplace_users, marketplace_subscriptions) new_id = model.entitlements.get_web_customer_ids(test_user.id) - assert new_id != [12345] + assert new_id != [55555] assert new_id == marketplace_users.lookup_customer_id(test_user.email) # make sure it will remove account numbers from db that do not belong