316 lines
12 KiB
Python
316 lines
12 KiB
Python
"""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_tracker_enables_billing_with_default_db_when_requested(tmp_path, monkeypatch):
|
|
from meshnet_tracker.billing import DEFAULT_BILLING_DB_PATH
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
tracker = TrackerServer(enable_billing=True)
|
|
port = tracker.start()
|
|
try:
|
|
summary = json.loads(
|
|
urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/billing/summary").read()
|
|
)
|
|
assert "protocol_cut" in summary
|
|
finally:
|
|
tracker.stop()
|
|
assert (tmp_path / DEFAULT_BILLING_DB_PATH).exists()
|
|
|
|
|
|
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)
|
|
|
|
stats = json.loads(urllib.request.urlopen(f"{tracker_url}/v1/stats").read())
|
|
node_stats = next(iter(stats["nodes"].values()))["models"][MODEL]
|
|
assert node_stats["tokens_per_sec_last_hour"] is not None
|
|
assert node_stats["sample_count_last_hour"] == 1
|
|
|
|
|
|
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)
|