billing ledger
This commit is contained in:
Dobromir Popov
294
tests/test_billing_ledger.py
Normal file
294
tests/test_billing_ledger.py
Normal file
@@ -0,0 +1,294 @@
|
||||
"""US-031: billing ledger — per-token pricing, 90/10 split, pending balances.
|
||||
|
||||
Unit tests for BillingLedger math/persistence/replication, plus HTTP
|
||||
integration: 401 without an API key, billed 200 with one, 402 once the
|
||||
Caller Credit is exhausted (rejected before routing — no free work).
|
||||
"""
|
||||
|
||||
import http.server
|
||||
import json
|
||||
import socketserver
|
||||
import threading
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
import pytest
|
||||
|
||||
from meshnet_tracker.billing import BillingLedger
|
||||
from meshnet_tracker.server import TrackerServer
|
||||
|
||||
MODEL = "openai-community/gpt2"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- unit tests
|
||||
|
||||
|
||||
def test_charge_single_node_gets_90_percent():
|
||||
ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
|
||||
ledger.ensure_client("key-a")
|
||||
event = ledger.charge_request("key-a", MODEL, total_tokens=1000, node_work=[("wallet-1", 12)])
|
||||
assert event["cost"] == pytest.approx(0.02)
|
||||
assert ledger.get_client_balance("key-a") == pytest.approx(1.0 - 0.02)
|
||||
assert ledger.get_node_pending("wallet-1") == pytest.approx(0.02 * 0.90)
|
||||
assert ledger.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10)
|
||||
|
||||
|
||||
def test_charge_three_node_split_by_work_units():
|
||||
ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
|
||||
ledger.ensure_client("key-a")
|
||||
ledger.charge_request(
|
||||
"key-a", MODEL, total_tokens=1000,
|
||||
node_work=[("wallet-1", 6), ("wallet-2", 3), ("wallet-3", 3)],
|
||||
)
|
||||
pool = 0.02 * 0.90
|
||||
assert ledger.get_node_pending("wallet-1") == pytest.approx(pool * 6 / 12)
|
||||
assert ledger.get_node_pending("wallet-2") == pytest.approx(pool * 3 / 12)
|
||||
assert ledger.get_node_pending("wallet-3") == pytest.approx(pool * 3 / 12)
|
||||
assert ledger.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10)
|
||||
|
||||
|
||||
def test_walletless_node_share_accrues_to_protocol_cut():
|
||||
ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
|
||||
ledger.ensure_client("key-a")
|
||||
ledger.charge_request(
|
||||
"key-a", MODEL, total_tokens=1000,
|
||||
node_work=[("wallet-1", 6), (None, 6)],
|
||||
)
|
||||
pool = 0.02 * 0.90
|
||||
assert ledger.get_node_pending("wallet-1") == pytest.approx(pool / 2)
|
||||
# walletless half of the pool + the 10% cut both land in protocol_cut
|
||||
assert ledger.snapshot()["protocol_cut"] == pytest.approx(pool / 2 + 0.02 * 0.10)
|
||||
|
||||
|
||||
def test_per_model_price_override():
|
||||
ledger = BillingLedger(
|
||||
starting_credit=1.0,
|
||||
default_price_per_1k=0.02,
|
||||
prices={MODEL: 0.10},
|
||||
)
|
||||
ledger.ensure_client("key-a")
|
||||
event = ledger.charge_request("key-a", MODEL, 500, [("wallet-1", 1)])
|
||||
assert event["cost"] == pytest.approx(0.10 * 500 / 1000)
|
||||
event = ledger.charge_request("key-a", "other-model", 500, [("wallet-1", 1)])
|
||||
assert event["cost"] == pytest.approx(0.02 * 500 / 1000)
|
||||
|
||||
|
||||
def test_payout_and_forfeit_hooks():
|
||||
ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
|
||||
ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 1)])
|
||||
pending = ledger.get_node_pending("wallet-1")
|
||||
ledger.settle_node_payout("wallet-1", pending, reference="tx-sig-1")
|
||||
assert ledger.get_node_pending("wallet-1") == pytest.approx(0.0)
|
||||
|
||||
ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 1)])
|
||||
cut_before = ledger.snapshot()["protocol_cut"]
|
||||
forfeited = ledger.forfeit_pending("wallet-1")["amount"]
|
||||
assert forfeited == pytest.approx(0.02 * 0.90)
|
||||
assert ledger.get_node_pending("wallet-1") == pytest.approx(0.0)
|
||||
assert ledger.snapshot()["protocol_cut"] == pytest.approx(cut_before + forfeited)
|
||||
|
||||
|
||||
def test_restart_persistence(tmp_path):
|
||||
db = str(tmp_path / "billing.db")
|
||||
ledger = BillingLedger(db_path=db, starting_credit=1.0, default_price_per_1k=0.02)
|
||||
ledger.credit_client("key-a", 5.0)
|
||||
ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 12)])
|
||||
ledger.save_to_db()
|
||||
|
||||
reloaded = BillingLedger(db_path=db, starting_credit=1.0, default_price_per_1k=0.02)
|
||||
assert reloaded.get_client_balance("key-a") == pytest.approx(
|
||||
ledger.get_client_balance("key-a")
|
||||
)
|
||||
assert reloaded.get_node_pending("wallet-1") == pytest.approx(0.02 * 0.90)
|
||||
assert reloaded.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10)
|
||||
|
||||
|
||||
def test_event_replication_converges_and_dedupes():
|
||||
leader = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
|
||||
follower = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02)
|
||||
|
||||
leader.credit_client("key-a", 10.0)
|
||||
leader.charge_request("key-a", MODEL, 2000, [("wallet-1", 8), ("wallet-2", 4)])
|
||||
|
||||
events, cursor = leader.events_since(0)
|
||||
assert follower.apply_events(events) == len(events)
|
||||
# replaying the same batch must be a no-op
|
||||
assert follower.apply_events(events) == 0
|
||||
|
||||
assert follower.get_client_balance("key-a") == pytest.approx(
|
||||
leader.get_client_balance("key-a")
|
||||
)
|
||||
assert follower.get_node_pending("wallet-1") == pytest.approx(
|
||||
leader.get_node_pending("wallet-1")
|
||||
)
|
||||
assert follower.snapshot()["protocol_cut"] == pytest.approx(
|
||||
leader.snapshot()["protocol_cut"]
|
||||
)
|
||||
# incremental cursor: nothing new after full sync
|
||||
more, _ = leader.events_since(cursor)
|
||||
assert more == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------- HTTP integration
|
||||
|
||||
|
||||
class _UsageStubNode:
|
||||
"""Minimal head node returning a chat completion with real usage numbers."""
|
||||
|
||||
def __init__(self, total_tokens: int = 1000):
|
||||
self.total_tokens = total_tokens
|
||||
outer = self
|
||||
|
||||
class Handler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, fmt, *args):
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
self.rfile.read(length)
|
||||
body = json.dumps({
|
||||
"id": "chatcmpl-stub",
|
||||
"object": "chat.completion",
|
||||
"model": MODEL,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "ok"},
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": outer.total_tokens,
|
||||
"total_tokens": outer.total_tokens,
|
||||
},
|
||||
}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
self._server = socketserver.ThreadingTCPServer(("127.0.0.1", 0), Handler)
|
||||
self._server.daemon_threads = True
|
||||
self._thread: threading.Thread | None = None
|
||||
|
||||
def start(self) -> int:
|
||||
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
||||
self._thread.start()
|
||||
return self._server.server_address[1]
|
||||
|
||||
def stop(self):
|
||||
self._server.shutdown()
|
||||
self._server.server_close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def billed_tracker():
|
||||
ledger = BillingLedger(starting_credit=0.03, default_price_per_1k=0.02)
|
||||
tracker = TrackerServer(
|
||||
model_presets={
|
||||
MODEL: {
|
||||
"layers_start": 0,
|
||||
"layers_end": 11,
|
||||
"bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
|
||||
}
|
||||
},
|
||||
billing=ledger,
|
||||
)
|
||||
port = tracker.start()
|
||||
tracker_url = f"http://127.0.0.1:{port}"
|
||||
|
||||
stub = _UsageStubNode(total_tokens=1000)
|
||||
stub_port = stub.start()
|
||||
reg = json.dumps({
|
||||
"endpoint": f"http://127.0.0.1:{stub_port}",
|
||||
"shard_start": 0,
|
||||
"shard_end": 11,
|
||||
"model": MODEL,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0,
|
||||
"tracker_mode": True,
|
||||
"wallet_address": "wallet-head",
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{tracker_url}/v1/nodes/register",
|
||||
data=reg,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req) as r:
|
||||
r.read()
|
||||
|
||||
yield tracker_url, ledger
|
||||
|
||||
stub.stop()
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def _chat(tracker_url: str, api_key: str | None):
|
||||
data = json.dumps({
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}).encode()
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
req = urllib.request.Request(
|
||||
f"{tracker_url}/v1/chat/completions",
|
||||
data=data,
|
||||
headers=headers,
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req) as r:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def test_proxy_chat_requires_api_key_when_billing_enabled(billed_tracker):
|
||||
tracker_url, _ = billed_tracker
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_chat(tracker_url, api_key=None)
|
||||
assert exc_info.value.code == 401
|
||||
|
||||
|
||||
def test_proxy_chat_bills_client_and_credits_node(billed_tracker):
|
||||
tracker_url, ledger = billed_tracker
|
||||
_chat(tracker_url, api_key="client-1")
|
||||
# 1000 tokens at 0.02/1K = 0.02 USDT; starting credit 0.03
|
||||
assert ledger.get_client_balance("client-1") == pytest.approx(0.03 - 0.02)
|
||||
assert ledger.get_node_pending("wallet-head") == pytest.approx(0.02 * 0.90)
|
||||
|
||||
summary = json.loads(
|
||||
urllib.request.urlopen(f"{tracker_url}/v1/billing/summary").read()
|
||||
)
|
||||
assert summary["node_pending"]["wallet-head"] == pytest.approx(0.02 * 0.90)
|
||||
assert summary["protocol_cut"] == pytest.approx(0.02 * 0.10)
|
||||
|
||||
|
||||
def test_proxy_chat_402_when_balance_exhausted(billed_tracker):
|
||||
tracker_url, ledger = billed_tracker
|
||||
_chat(tracker_url, api_key="client-2") # 0.03 -> 0.01
|
||||
_chat(tracker_url, api_key="client-2") # 0.01 -> -0.01 (post-pay drift)
|
||||
assert ledger.get_client_balance("client-2") == pytest.approx(-0.01)
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
_chat(tracker_url, api_key="client-2")
|
||||
assert exc_info.value.code == 402
|
||||
# rejected before routing: nothing further was billed
|
||||
assert ledger.get_client_balance("client-2") == pytest.approx(-0.01)
|
||||
|
||||
|
||||
def test_billing_gossip_endpoint_applies_events(billed_tracker):
|
||||
tracker_url, ledger = billed_tracker
|
||||
peer = BillingLedger(starting_credit=0.03, default_price_per_1k=0.02)
|
||||
peer.credit_client("remote-client", 7.0)
|
||||
events, _ = peer.events_since(0)
|
||||
body = json.dumps({"events": events}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{tracker_url}/v1/billing/gossip",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req) as r:
|
||||
assert json.loads(r.read())["applied"] == len(events)
|
||||
# only the replicated credit event lands here — Caller Credit was granted
|
||||
# on the peer that first saw the key, and its event replicates separately
|
||||
assert ledger.get_client_balance("remote-client") == pytest.approx(7.0)
|
||||
Reference in New Issue
Block a user