apunkt/dograh
1
0
Fork 0
mirror of https://github.com/dograh-hq/dograh.git synced 2026-07-10 11:12:13 +02:00
Code Issues Projects Releases Packages Wiki Activity

chore: cleaup mps v1 billing (#507)

Browse source 
* chore: cleaup mps v1 billing

* chore: remove legacy file upload path

* chore: implement review comments
This commit is contained in:
Abhishek 2026-07-07 18:38:29 +05:30 • committed by GitHub
parent ac01f7775e
commit fdb7f92fcc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
36 changed files with 268 additions and 1319 deletions
Download patch file Download diff file Expand all files Collapse all files

9
api/tests/test_ai_model_configuration_v2.py
View file

@ -1,5 +1,5 @@
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
from unittest.mock import AsyncMock
import pytest
from pydantic import ValidationError
@ -442,7 +442,6 @@ async def test_migrate_model_configuration_v2_initializes_hosted_mps_billing(
ensure_billing = AsyncMock(return_value={"billing_mode": "v2"})
upsert = AsyncMock()
migrate_workflows = AsyncMock()
sync_posthog_billing = Mock()
monkeypatch.setattr(organization_routes, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
@ -480,11 +479,6 @@ async def test_migrate_model_configuration_v2_initializes_hosted_mps_billing(
"_model_configuration_v2_response",
AsyncMock(return_value=expected_response),
)
monkeypatch.setattr(
organization_routes,
"_sync_posthog_organization_mps_billing_v2_status",
sync_posthog_billing,
)
user = SimpleNamespace(
id=7,
@ -503,5 +497,4 @@ async def test_migrate_model_configuration_v2_initializes_hosted_mps_billing(
organization_id=42,
fallback_user_config=legacy,
)
sync_posthog_billing.assert_called_once_with(42, uses_mps_billing_v2=True)
assert response == expected_response

66
api/tests/test_auth_depends.py
View file

@ -209,7 +209,7 @@ def test_associate_user_with_posthog_org_supports_backfill_arguments(monkeypatch
assert "backfilled" not in capture_kwargs["properties"]
def test_sync_created_organization_to_posthog_supports_billing_flag(monkeypatch):
def test_sync_created_organization_to_posthog_with_provider_id(monkeypatch):
organization = SimpleNamespace(id=42, provider_id="team-1")
group_calls = []
capture_calls = []
@ -228,11 +228,9 @@ def test_sync_created_organization_to_posthog_supports_billing_flag(monkeypatch)
auth_depends._sync_created_organization_to_posthog(
organization=organization,
created_by_provider_id="stack-user-1",
uses_mps_billing_v2=True,
)
_, group_kwargs = group_calls[0]
group_args, _ = group_calls[0]
group_args, group_kwargs = group_calls[0]
assert group_args == (
"organization",
"42",
@ -241,69 +239,9 @@ def test_sync_created_organization_to_posthog_supports_billing_flag(monkeypatch)
"organization_provider_id": "team-1",
"auth_provider": "stack",
"created_by_provider_id": "stack-user-1",
"uses_mps_billing_v2": True,
},
)
assert group_kwargs == {"distinct_id": "stack-user-1"}
_, capture_kwargs = capture_calls[0]
assert capture_kwargs["distinct_id"] == "stack-user-1"
assert capture_kwargs["properties"]["uses_mps_billing_v2"] is True
def test_sync_posthog_organization_group_properties_has_no_distinct_id(monkeypatch):
organization = SimpleNamespace(id=42, provider_id="team-1")
group_calls = []
monkeypatch.setattr(
auth_depends,
"group_identify",
lambda *args, **kwargs: group_calls.append((args, kwargs)),
)
auth_depends._sync_posthog_organization_group_properties(
organization=organization,
uses_mps_billing_v2=True,
)
assert group_calls == [
(
(
"organization",
"42",
{
"organization_id": 42,
"organization_provider_id": "team-1",
"auth_provider": "stack",
"uses_mps_billing_v2": True,
},
),
{},
)
]
def test_sync_posthog_organization_mps_billing_v2_status(monkeypatch):
group_calls = []
monkeypatch.setattr(
auth_depends,
"group_identify",
lambda *args, **kwargs: group_calls.append((args, kwargs)),
)
auth_depends._sync_posthog_organization_mps_billing_v2_status(
42,
uses_mps_billing_v2=True,
)
assert group_calls == [
(
(
"organization",
"42",
{"uses_mps_billing_v2": True},
),
{},
)
]

80
api/tests/test_dograh_embedding_service.py
View file

@ -49,114 +49,52 @@ async def test_dograh_embedding_sends_plain_without_correlation():
await service.embed_texts(["hello"])
create.assert_awaited_once()
# No correlation id (e.g. a v1 org) → no MPS metadata; MPS accepts plain calls.
# No correlation id → no MPS metadata; MPS accepts plain calls.
assert "extra_body" not in create.await_args.kwargs
def _fake_mps_client(*, status_return=None, minted="minted"):
def _fake_mps_client(*, minted="minted"):
return SimpleNamespace(
get_billing_account_status=AsyncMock(return_value=status_return),
create_correlation_id=AsyncMock(return_value={"correlation_id": minted}),
)
@pytest.mark.asyncio
async def test_resolve_correlation_oss_mints_directly(monkeypatch):
async def test_resolve_correlation_mints_via_service_key(monkeypatch):
fake = _fake_mps_client()
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "oss")
result = await resolve_embedding_correlation_id(
organization_id=None, service_key="sk-mps"
)
result = await resolve_embedding_correlation_id(service_key="sk-mps")
assert result == "minted"
fake.create_correlation_id.assert_awaited_once_with(service_key="sk-mps")
fake.get_billing_account_status.assert_not_awaited()
@pytest.mark.asyncio
async def test_resolve_correlation_hosted_v2_mints(monkeypatch):
fake = _fake_mps_client(status_return={"billing_mode": "v2"})
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "hosted")
result = await resolve_embedding_correlation_id(
organization_id=42, service_key="sk-mps", created_by="user-1"
)
assert result == "minted"
fake.get_billing_account_status.assert_awaited_once_with(42, created_by="user-1")
fake.create_correlation_id.assert_awaited_once_with(service_key="sk-mps")
@pytest.mark.asyncio
async def test_resolve_correlation_hosted_v1_returns_none_without_minting(monkeypatch):
fake = _fake_mps_client(status_return={"billing_mode": "v1"})
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "hosted")
result = await resolve_embedding_correlation_id(
organization_id=42, service_key="sk-mps"
)
assert result is None
fake.create_correlation_id.assert_not_awaited()
@pytest.mark.asyncio
async def test_resolve_correlation_hosted_no_account_returns_none(monkeypatch):
fake = _fake_mps_client(status_return=None)
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "hosted")
result = await resolve_embedding_correlation_id(
organization_id=42, service_key="sk-mps"
)
assert result is None
fake.create_correlation_id.assert_not_awaited()
@pytest.mark.asyncio
async def test_resolve_correlation_no_service_key_returns_none(monkeypatch):
fake = _fake_mps_client(status_return={"billing_mode": "v2"})
fake = _fake_mps_client()
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "hosted")
result = await resolve_embedding_correlation_id(
organization_id=42, service_key=None
)
result = await resolve_embedding_correlation_id(service_key=None)
assert result is None
fake.get_billing_account_status.assert_not_awaited()
fake.create_correlation_id.assert_not_awaited()
@pytest.mark.asyncio
async def test_resolve_correlation_swallows_errors(monkeypatch):
fake = SimpleNamespace(
get_billing_account_status=AsyncMock(side_effect=RuntimeError("mps down")),
create_correlation_id=AsyncMock(),
create_correlation_id=AsyncMock(side_effect=RuntimeError("mps down")),
)
monkeypatch.setattr(
"api.services.mps_service_key_client.mps_service_key_client", fake
)
monkeypatch.setattr("api.constants.DEPLOYMENT_MODE", "hosted")
# A transient MPS failure must not break embeddings — fall back to no protocol.
result = await resolve_embedding_correlation_id(
organization_id=42, service_key="sk-mps"
)
# A transient MPS failure must not break embeddings — send without it.
result = await resolve_embedding_correlation_id(service_key="sk-mps")
assert result is None

45
api/tests/test_mps_service_key_client.py
View file

@ -130,51 +130,6 @@ async def test_create_correlation_id_uses_bearer_auth(monkeypatch):
]
@pytest.mark.asyncio
async def test_get_billing_account_status_uses_hosted_org_auth(monkeypatch):
calls = []
class FakeAsyncClient:
def __init__(self, timeout):
self.timeout = timeout
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return None
async def get(self, url, headers):
calls.append(("GET", url, headers))
return _Response(200, {"organization_id": 42, "billing_mode": "v2"})
monkeypatch.setattr(
"api.services.mps_service_key_client.httpx.AsyncClient", FakeAsyncClient
)
monkeypatch.setattr("api.services.mps_service_key_client.DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
"api.services.mps_service_key_client.DOGRAH_MPS_SECRET_KEY", "mps-secret"
)
client = MPSServiceKeyClient()
assert await client.get_billing_account_status(organization_id=42) == {
"organization_id": 42,
"billing_mode": "v2",
}
assert calls == [
(
"GET",
f"{client.base_url}/api/v1/billing/accounts/42/status",
{
"Content-Type": "application/json",
"X-Secret-Key": "mps-secret",
"X-Organization-Id": "42",
},
)
]
@pytest.mark.asyncio
async def test_authorize_workflow_run_start_uses_hosted_org_auth(monkeypatch):
calls = []

44
api/tests/test_organization_usage_billing.py
View file

@ -6,41 +6,36 @@ import pytest
from api.routes import organization_usage
def test_is_mps_billing_v2_depends_only_on_account_mode():
assert organization_usage._is_mps_billing_v2({"billing_mode": "v2"}) is True
assert organization_usage._is_mps_billing_v2({"billing_mode": "v1"}) is False
assert organization_usage._is_mps_billing_v2({"billing_mode": "shadow"}) is False
assert organization_usage._is_mps_billing_v2(None) is False
@pytest.mark.asyncio
async def test_get_mps_billing_account_status_uses_user_provider_id(monkeypatch):
get_status = AsyncMock(return_value={"billing_mode": "v2"})
async def test_get_billing_credits_oss_aggregates_by_created_by(monkeypatch):
monkeypatch.setattr(organization_usage, "DEPLOYMENT_MODE", "oss")
get_usage = AsyncMock(
return_value={"total_credits_used": 12.5, "remaining_credits": 487.5}
)
monkeypatch.setattr(
organization_usage.mps_service_key_client,
"get_billing_account_status",
get_status,
"get_usage_by_created_by",
get_usage,
)
user = SimpleNamespace(provider_id="provider-123")
user = SimpleNamespace(provider_id="provider-123", selected_organization_id=None)
assert await organization_usage._get_mps_billing_account_status(user, 42) == {
"billing_mode": "v2"
}
get_status.assert_awaited_once_with(
organization_id=42,
created_by="provider-123",
response = await organization_usage.get_billing_credits(
page=1,
limit=50,
user=user,
)
get_usage.assert_awaited_once_with("provider-123")
assert response.total_credits_used == 12.5
assert response.remaining_credits == 487.5
assert response.total_quota == 500.0
assert response.ledger_entries == []
@pytest.mark.asyncio
async def test_get_billing_credits_pages_v2_ledger(monkeypatch):
async def test_get_billing_credits_pages_hosted_ledger(monkeypatch):
monkeypatch.setattr(organization_usage, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
organization_usage,
"_get_mps_billing_account_status",
AsyncMock(return_value={"billing_mode": "v2"}),
)
get_ledger = AsyncMock(
return_value={
"account": {
@ -90,7 +85,6 @@ async def test_get_billing_credits_pages_v2_ledger(monkeypatch):
limit=25,
created_by="provider-123",
)
assert response.billing_version == "v2"
assert response.total_credits_used == 75
assert response.total_count == 101
assert response.page == 3

29
api/tests/test_quota_service.py
View file

@ -166,34 +166,22 @@ async def test_authorize_workflow_run_v2_insufficient_credits_prompts_billing(
@pytest.mark.asyncio
async def test_authorize_workflow_run_v1_uses_legacy_key_usage(
async def test_authorize_workflow_run_oss_exhausted_key_blocks_run(
monkeypatch,
):
api_key = "mps_sk_12345678"
get_config = AsyncMock(return_value=_dograh_config(api_key))
authorize = AsyncMock(
return_value={
"allowed": True,
"billing_mode": "v1",
"remaining_credits": "0.0000",
}
)
check_usage = AsyncMock(
return_value={"total_credits_used": 500.0, "remaining_credits": 0.0}
)
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(quota_service, "DEPLOYMENT_MODE", "oss")
_patch_workflow_context(monkeypatch)
monkeypatch.setattr(
quota_service,
"get_effective_ai_model_configuration_for_workflow",
get_config,
)
monkeypatch.setattr(
quota_service.mps_service_key_client,
"authorize_workflow_run_start",
authorize,
)
monkeypatch.setattr(
quota_service.mps_service_key_client,
"check_service_key_usage",
@ -206,12 +194,7 @@ async def test_authorize_workflow_run_v1_uses_legacy_key_usage(
assert result.error_code == "quota_exceeded"
assert "founders@dograh.com" in result.error_message
assert "/billing" not in result.error_message
authorize.assert_awaited_once()
check_usage.assert_awaited_once_with(
api_key,
organization_id=42,
created_by="provider-123",
)
check_usage.assert_awaited_once_with(api_key)
@pytest.mark.asyncio
@ -404,11 +387,7 @@ async def test_authorize_workflow_run_oss_uses_key_paths_not_workflow_org(
assert result.has_quota is True
hosted_authorize.assert_not_awaited()
check_usage.assert_awaited_once_with(
api_key,
organization_id=None,
created_by="provider-123",
)
check_usage.assert_awaited_once_with(api_key)
create_correlation.assert_awaited_once_with(
service_key=api_key,
workflow_run_id=88,

4
api/tests/test_workflow_graph_constraints.py
View file

@ -9,9 +9,9 @@ category of violation we found in production. We pin two layers:
layer ever stops rejecting one of these fixtures, the production
write paths will quietly start accepting bad workflows again.
2. audit_definition (api.services.workflow.audit) — read-only sweep
over persisted rows used by the admin cleanup script to find
over persisted rows for one-off cleanup tooling that finds
legacy/imported breakage. Pinned so refactors of the rule set
don't silently change the verdicts the migration relies on.
don't silently change those cleanup verdicts.
DTO-level shape validation is covered by `test_dto.py` and isn't
re-pinned here.

49
api/tests/test_workflow_run_billing.py
View file

@ -40,15 +40,9 @@ async def test_report_workflow_run_platform_usage_reports_hosted_completion(
monkeypatch,
):
workflow_run = _make_workflow_run()
get_status = AsyncMock(return_value={"billing_mode": "v2"})
report_usage = AsyncMock(return_value={"metered": True})
monkeypatch.setattr(workflow_run_billing_mod, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"get_billing_account_status",
get_status,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"report_platform_usage",
@ -76,15 +70,9 @@ async def test_report_workflow_run_platform_usage_reports_duration_without_corre
):
workflow_run = _make_workflow_run()
workflow_run.initial_context = {}
get_status = AsyncMock(return_value={"billing_mode": "v2"})
report_usage = AsyncMock(return_value={"metered": True})
monkeypatch.setattr(workflow_run_billing_mod, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"get_billing_account_status",
get_status,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"report_platform_usage",
@ -106,30 +94,6 @@ async def test_report_workflow_run_platform_usage_reports_duration_without_corre
)
@pytest.mark.asyncio
async def test_report_workflow_run_platform_usage_skips_non_v2_account(monkeypatch):
workflow_run = _make_workflow_run()
get_status = AsyncMock(return_value={"billing_mode": "v1"})
report_usage = AsyncMock()
monkeypatch.setattr(workflow_run_billing_mod, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"get_billing_account_status",
get_status,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"report_platform_usage",
report_usage,
)
await report_workflow_run_platform_usage(workflow_run)
get_status.assert_awaited_once_with(organization_id=42)
report_usage.assert_not_awaited()
@pytest.mark.asyncio
async def test_report_workflow_run_platform_usage_skips_missing_duration_without_correlation(
monkeypatch,
@ -137,15 +101,9 @@ async def test_report_workflow_run_platform_usage_skips_missing_duration_without
workflow_run = _make_workflow_run()
workflow_run.initial_context = {}
workflow_run.usage_info = {}
get_status = AsyncMock(return_value={"billing_mode": "v2"})
report_usage = AsyncMock()
monkeypatch.setattr(workflow_run_billing_mod, "DEPLOYMENT_MODE", "saas")
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"get_billing_account_status",
get_status,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"report_platform_usage",
@ -154,7 +112,6 @@ async def test_report_workflow_run_platform_usage_skips_missing_duration_without
await report_workflow_run_platform_usage(workflow_run)
get_status.assert_not_awaited()
report_usage.assert_not_awaited()
@ -197,7 +154,6 @@ async def test_report_workflow_run_platform_usage_skips_incomplete(monkeypatch):
async def test_report_completed_workflow_run_platform_usage_loads_run(monkeypatch):
workflow_run = _make_workflow_run()
get_run = AsyncMock(return_value=workflow_run)
get_status = AsyncMock(return_value={"billing_mode": "v2"})
report_usage = AsyncMock(return_value={"metered": True})
monkeypatch.setattr(workflow_run_billing_mod, "DEPLOYMENT_MODE", "saas")
@ -206,11 +162,6 @@ async def test_report_completed_workflow_run_platform_usage_loads_run(monkeypatc
"get_workflow_run_by_id",
get_run,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"get_billing_account_status",
get_status,
)
monkeypatch.setattr(
workflow_run_billing_mod.mps_service_key_client,
"report_platform_usage",