Skip to content

Commit

Permalink
feat: Update tiers usage information with local org usage information (…
Browse files Browse the repository at this point in the history 
…#21957)

* update tiers usage with local org usage information

* type finagling

* update addon tiers as well

* test

* Update ee/billing/billing_manager.py

* update tests

---------

Co-authored-by: Raquel Smith <[email protected]>
  • Loading branch information
@xrdt @raquelmsmith
xrdt and raquelmsmith authored May 7, 2024
1 parent f1d32e6 commit a266ffc
Show file tree
Hide file tree
Showing 5 changed files with 355 additions and 30 deletions.
2 changes: 1 addition & 1 deletion .vscode/launch.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@
"DATABASE_URL": "postgres://posthog:posthog@localhost:5432/posthog",
"SKIP_SERVICE_VERSION_REQUIREMENTS": "1",
"PRINT_SQL": "1",
"BILLING_SERVICE_URL": "https://billing.dev.posthog.dev"
"BILLING_SERVICE_URL": "http://localhost:8100" //"https://billing.dev.posthog.dev"
},
"console": "integratedTerminal",
"python": "${workspaceFolder}/env/bin/python",
Expand Down
236 changes: 222 additions & 14 deletions ee/api/test/test_billing.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,12 @@
from rest_framework import status

from ee.api.test.base import APILicensedTest
from ee.billing.billing_types import BillingPeriod, CustomerInfo, CustomerProduct
from ee.billing.billing_types import (
BillingPeriod,
CustomerInfo,
CustomerProduct,
CustomerProductAddon,
)
from ee.models.license import License
from posthog.cloud_utils import (
TEST_clear_instance_license_cache,
Expand Down Expand Up @@ -76,7 +81,7 @@ def create_billing_customer(**kwargs) -> CustomerInfo:
{
"unit_amount_usd": "0.00045",
"up_to": 2000000,
"current_amount_usd": None,
"current_amount_usd": "0.00",
},
],
tiered=True,
Expand All @@ -89,6 +94,39 @@ def create_billing_customer(**kwargs) -> CustomerInfo:
projected_usage=0,
projected_amount_usd="0.00",
usage_key="events",
addons=[
CustomerProductAddon(
name="Addon",
description="Test Addon",
price_description=None,
type="events",
image_url="https://posthog.com/static/images/product-os.png",
free_allocation=10000,
tiers=[
{
"unit_amount_usd": "0.00",
"up_to": 1000000,
"current_amount_usd": "0.00",
},
{
"unit_amount_usd": "0.0000135",
"up_to": 2000000,
"current_amount_usd": "0.00",
},
],
tiered=True,
unit_amount_usd="0.00",
current_amount_usd="0.00",
current_usage=0,
usage_limit=None,
has_exceeded_limit=False,
percentage_usage=0,
projected_usage=0,
projected_amount_usd="0.00",
usage_key="events",
subscribed=True,
)
],
)
],
billing_period=BillingPeriod(
Expand Down Expand Up @@ -121,22 +159,72 @@ def create_billing_products_response(**kwargs) -> dict[str, list[CustomerProduct
"unit_amount_usd": "0.00",
"up_to": 1000000,
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
},
{
"unit_amount_usd": "0.00045",
"up_to": 2000000,
"current_amount_usd": None,
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
},
],
addons=[
{
"current_amount_usd": 0.0,
"current_usage": 0,
"description": "Test Addon",
"free_allocation": 10000,
"has_exceeded_limit": False,
"image_url": "https://posthog.com/static/images/product-os.png",
"name": "Addon",
"percentage_usage": 0,
"price_description": None,
"projected_amount_usd": "0.00",
"projected_usage": 0,
"subscribed": True,
"tiered": True,
"tiers": [
{
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
"unit_amount_usd": "0.00",
"up_to": 1000000,
},
{
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
"unit_amount_usd": "0.0000135",
"up_to": 2000000,
},
],
"type": "events",
"unit_amount_usd": "0.00",
"usage_key": "events",
"usage_limit": None,
},
],
tiered=True,
unit_amount_usd="0.00",
current_amount_usd="0.00",
current_amount_usd=0.0,
current_usage=0,
usage_limit=None,
has_exceeded_limit=False,
percentage_usage=0,
projected_usage=0,
projected_amount_usd="0.00",
projected_amount=0,
projected_amount_usd=0.00,
usage_key="events",
)
]
Expand Down Expand Up @@ -271,15 +359,23 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma
"unit_amount_usd": "0.00",
"up_to": 1000000,
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
},
{
"unit_amount_usd": "0.00045",
"up_to": 2000000,
"current_amount_usd": None,
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
},
],
"tiered": True,
"current_amount_usd": "0.00",
"current_amount_usd": 0.00,
"current_usage": 0,
"usage_limit": None,
"percentage_usage": 0,
Expand All @@ -288,7 +384,48 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma
"projected_amount_usd": "0.00",
"projected_usage": 0,
"usage_key": "events",
}
"addons": [
{
"current_amount_usd": 0.00,
"current_usage": 0,
"description": "Test Addon",
"free_allocation": 10000,
"has_exceeded_limit": False,
"image_url": "https://posthog.com/static/images/product-os.png",
"name": "Addon",
"percentage_usage": 0,
"price_description": None,
"projected_amount_usd": "0.00",
"projected_usage": 0,
"subscribed": True,
"tiered": True,
"tiers": [
{
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
"unit_amount_usd": "0.00",
"up_to": 1000000,
},
{
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
"unit_amount_usd": "0.0000135",
"up_to": 2000000,
},
],
"type": "events",
"unit_amount_usd": "0.00",
"usage_key": "events",
"usage_limit": None,
},
],
},
],
"billing_period": {
"current_period_start": "2022-10-07T11:12:48",
Expand Down Expand Up @@ -342,25 +479,75 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma
"unit_amount_usd": "0.00",
"up_to": 1000000,
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
},
{
"unit_amount_usd": "0.00045",
"up_to": 2000000,
"current_amount_usd": None,
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
},
],
"current_usage": 0,
"percentage_usage": 0.0,
"current_amount_usd": "0.00",
"percentage_usage": 0,
"current_amount_usd": 0.0,
"has_exceeded_limit": False,
"projected_amount_usd": "0.00",
"projected_amount_usd": 0.0,
"projected_amount": 0,
"projected_usage": 0,
"tiered": True,
"unit_amount_usd": "0.00",
"usage_limit": None,
"image_url": "https://posthog.com/static/images/product-os.png",
"percentage_usage": 0.0,
"percentage_usage": 0,
"usage_key": "events",
"addons": [
{
"current_amount_usd": 0.0,
"current_usage": 0,
"description": "Test Addon",
"free_allocation": 10000,
"has_exceeded_limit": False,
"image_url": "https://posthog.com/static/images/product-os.png",
"name": "Addon",
"percentage_usage": 0,
"price_description": None,
"projected_amount_usd": "0.00",
"projected_usage": 0,
"subscribed": True,
"tiered": True,
"tiers": [
{
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
"unit_amount_usd": "0.00",
"up_to": 1000000,
},
{
"current_amount_usd": "0.00",
"current_usage": 0,
"flat_amount_usd": "0",
"projected_amount_usd": "None",
"projected_usage": None,
"unit_amount_usd": "0.0000135",
"up_to": 2000000,
},
],
"type": "events",
"unit_amount_usd": "0.00",
"usage_key": "events",
"usage_limit": None,
},
],
}
],
"billing_period": {
Expand Down Expand Up @@ -502,7 +689,7 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma
elif "api/billing" in url:
mock.status_code = 200
mock.json.return_value = create_billing_response(
customer=create_billing_customer(has_active_subscription=True)
customer=create_billing_customer(has_active_subscription=True),
)
mock.json.return_value["customer"]["usage_summary"]["events"]["usage"] = 1000

Expand Down Expand Up @@ -537,6 +724,27 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma
"period": ["2022-10-07T11:12:48", "2022-11-07T11:12:48"],
}

self.organization.usage = {"events": {"limit": None, "usage": 1000, "todays_usage": 1100000}}
self.organization.save()

res = self.client.get("/api/billing-v2")
assert res.status_code == 200
res_json = res.json()
# Should update product usage to reflect today's usage
assert res_json["products"][0]["current_usage"] == 1101000
assert res_json["products"][0]["current_amount_usd"] == 45.45
assert res_json["products"][0]["tiers"][0]["current_usage"] == 1000000
assert res_json["products"][0]["tiers"][0]["current_amount_usd"] == "0.00"
assert res_json["products"][0]["tiers"][1]["current_usage"] == 101000
assert res_json["products"][0]["tiers"][1]["current_amount_usd"] == "45.45"

assert res_json["products"][0]["addons"][0]["current_usage"] == 1101000
assert res_json["products"][0]["addons"][0]["current_amount_usd"] == 1.36
assert res_json["products"][0]["addons"][0]["tiers"][0]["current_usage"] == 1000000
assert res_json["products"][0]["addons"][0]["tiers"][0]["current_amount_usd"] == "0.00"
assert res_json["products"][0]["addons"][0]["tiers"][1]["current_usage"] == 101000
assert res_json["products"][0]["addons"][0]["tiers"][1]["current_amount_usd"] == "1.36"

def mock_implementation_missing_customer(url: str, headers: Any = None, params: Any = None) -> MagicMock:
mock = MagicMock()
mock.status_code = 404
Expand Down
Loading

1 comment on commit a266ffc

@feedanal
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This commit broke all self-hosted instances:

posthog-web-1                     | 2024-05-10T08:59:55.509759Z [error    ] 'subscribed'                   [posthog.exceptions] 
posthog-web-1                     | Traceback (most recent call last):
posthog-web-1                     |   File "/python-runtime/rest_framework/views.py", line 506, in dispatch
posthog-web-1                     |     response = handler(request, *args, **kwargs)
posthog-web-1                     |   File "/code/ee/api/billing.py", line 56, in list
posthog-web-1                     |     response = BillingManager(license).get_billing(org, plan_keys)
posthog-web-1                     |   File "/code/ee/billing/billing_manager.py", line 177, in get_billing
posthog-web-1                     |     if not addon["subscribed"]:
posthog-web-1                     | KeyError: 'subscribed'
posthog-web-1                     | 2024-05-10T08:59:55.603212Z [info     ] request_finished               [django_structlog.middlewares.request] code=500 
posthog-web-1                     | 2024-05-10T08:59:55.603406Z [error    ] Internal Server Error: /api/billing-v2/ [django.request] 
posthog-web-1                     | 2024-05-10T08:59:55.603434Z [error    ] Internal Server Error: /api/billing-v2/ [django.request] 
posthog-plugins-1                 | {"level":"info","time":1715331595667,"pid":144,"hostname":"9cdfa32006f1","groupId":"scheduled-tasks-runner","offsets":{},"msg":"[MAIN] ℹ️ consumer_status"}
posthog-web-1                     | 2024-05-10T08:59:55.987019Z [info     ] request_started

Please sign in to comment.