mirror of
https://github.com/dograh-hq/dograh.git
synced 2026-06-10 08:05:22 +02:00
87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
from datetime import UTC, datetime
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from api.services.pricing import workflow_run_cost as workflow_run_cost_mod
|
|
from api.services.pricing.workflow_run_cost import (
|
|
build_workflow_run_cost_info,
|
|
calculate_workflow_run_cost,
|
|
)
|
|
|
|
|
|
def _make_workflow_run():
|
|
return SimpleNamespace(
|
|
id=123,
|
|
workflow_id=456,
|
|
mode="textchat",
|
|
created_at=datetime.now(UTC),
|
|
usage_info={
|
|
"llm": {},
|
|
"tts": {},
|
|
"stt": {},
|
|
"call_duration_seconds": 7,
|
|
},
|
|
cost_info={},
|
|
workflow=SimpleNamespace(
|
|
organization_id=42,
|
|
user=SimpleNamespace(selected_organization_id=42),
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build_workflow_run_cost_info_does_not_update_org_usage(monkeypatch):
|
|
workflow_run = _make_workflow_run()
|
|
get_org = AsyncMock(return_value=SimpleNamespace(id=42, price_per_second_usd=1.5))
|
|
update_usage = AsyncMock()
|
|
|
|
monkeypatch.setattr(
|
|
workflow_run_cost_mod.db_client, "get_organization_by_id", get_org
|
|
)
|
|
monkeypatch.setattr(
|
|
workflow_run_cost_mod.db_client, "update_usage_after_run", update_usage
|
|
)
|
|
|
|
cost_info = await build_workflow_run_cost_info(workflow_run)
|
|
|
|
assert cost_info is not None
|
|
assert cost_info["call_duration_seconds"] == 7
|
|
assert "cost_breakdown" in cost_info
|
|
assert "dograh_token_usage" in cost_info
|
|
assert cost_info["charge_usd"] == 10.5
|
|
update_usage.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_calculate_workflow_run_cost_keeps_org_usage_side_effect_in_wrapper(
|
|
monkeypatch,
|
|
):
|
|
workflow_run = _make_workflow_run()
|
|
get_org = AsyncMock(return_value=SimpleNamespace(id=42, price_per_second_usd=None))
|
|
update_run = AsyncMock()
|
|
update_usage = AsyncMock()
|
|
|
|
monkeypatch.setattr(
|
|
workflow_run_cost_mod.db_client,
|
|
"get_workflow_run_by_id",
|
|
AsyncMock(return_value=workflow_run),
|
|
)
|
|
monkeypatch.setattr(
|
|
workflow_run_cost_mod.db_client, "get_organization_by_id", get_org
|
|
)
|
|
monkeypatch.setattr(
|
|
workflow_run_cost_mod.db_client, "update_workflow_run", update_run
|
|
)
|
|
monkeypatch.setattr(
|
|
workflow_run_cost_mod.db_client, "update_usage_after_run", update_usage
|
|
)
|
|
|
|
await calculate_workflow_run_cost(workflow_run.id)
|
|
|
|
update_run.assert_awaited_once()
|
|
saved_kwargs = update_run.await_args.kwargs
|
|
assert saved_kwargs["run_id"] == workflow_run.id
|
|
assert "cost_breakdown" in saved_kwargs["cost_info"]
|
|
update_usage.assert_awaited_once()
|