Skip to content

Commit

Permalink
free sku integration for reconciliation worker
Browse files Browse the repository at this point in the history
  • Loading branch information
Marcusk19 committed Feb 4, 2025
1 parent 37d65f7 commit ce2549e
Show file tree
Hide file tree
Showing 5 changed files with 171 additions and 36 deletions.
8 changes: 8 additions & 0 deletions data/billing.py
Original file line number Diff line number Diff line change
Expand Up @@ -372,6 +372,14 @@
"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]
Expand Down
78 changes: 76 additions & 2 deletions util/marketplace.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ def lookup_subscription(self, webCustomerId, skuId):
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):
Expand Down Expand Up @@ -171,6 +171,7 @@ def create_entitlement(self, customerId, sku):
"webCustomerId": customerId,
}
logger.debug("Created entitlement")

try:
r = requests.request(
method="post",
Expand All @@ -187,6 +188,30 @@ def create_entitlement(self, customerId, sku):

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
Expand Down Expand Up @@ -269,7 +294,6 @@ def get_list_of_subscriptions(

# Mocked classes for unit tests


TEST_USER = {
"account_number": 12345,
"email": "[email protected]",
Expand Down Expand Up @@ -358,6 +382,45 @@ def get_list_of_subscriptions(
"email": "[email protected]",
"username": "free_user",
}
PAID_USER = {
"account_number": 34567,
"email": "[email protected]",
"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):
Expand All @@ -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"]:
Expand All @@ -391,11 +456,20 @@ def lookup_subscription(self, customer_id, sku_id):
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

Expand Down
2 changes: 2 additions & 0 deletions util/test/test_marketplace.py
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,8 @@ def test_timeout_exception(self, requests_mock):
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):
Expand Down
102 changes: 71 additions & 31 deletions workers/reconciliationworker.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
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
Expand All @@ -24,6 +24,8 @@
SECONDS_IN_DAYS = 86400
ONE_MONTH = 30 * SECONDS_IN_DAYS * MILLISECONDS_IN_SECONDS

FREE_TIER_SKU = "MW04192"


class ReconciliationWorker(Worker):
def __init__(self):
Expand All @@ -42,9 +44,9 @@ def _perform_reconciliation(self, user_api, marketplace_api):

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)
Expand Down Expand Up @@ -77,39 +79,77 @@ def _perform_reconciliation(self, user_api, marketplace_api):
for customer_id in model_customer_ids:
if customer_id not in customer_ids:
entitlements.remove_web_customer_id(user, customer_id)

# 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)
# 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

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
Expand Down
17 changes: 14 additions & 3 deletions workers/test/test_reconciliationworker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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", "[email protected]")
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):
Expand Down Expand Up @@ -77,12 +88,12 @@ def test_reconcile_different_ids(initialized_db):
test_user = model.user.create_user("stripe_user", "password", "[email protected]")
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
Expand Down

0 comments on commit ce2549e

Please sign in to comment.