neuron-tai/tests/test_accounts.py

"""Dashboard user accounts: registration, login, roles, API keys, usage.

Unit tests for AccountStore plus HTTP integration on the tracker:
register/login/logout, per-account balance and usage, API-key lifecycle
(revoked keys rejected by the OpenAI proxy), and the admin listing.
"""

import json
import urllib.error
import urllib.request

import pytest

from meshnet_tracker.accounts import AccountStore
from meshnet_tracker.auth import sign_hive_request
from meshnet_tracker.billing import BillingLedger
from meshnet_tracker.server import TrackerServer

HIVE_SECRET = "test-hive-secret"


# ---------------------------------------------------------------- unit tests


def test_first_account_is_admin_then_users():
    store = AccountStore()
    first = store.register(email="admin@example.com", password="secret-123")
    second = store.register(email="user@example.com", password="secret-123")
    assert first["role"] == "admin"
    assert second["role"] == "user"


def test_register_requires_email_or_wallet_and_password_length():
    store = AccountStore()
    with pytest.raises(ValueError, match="email or a wallet"):
        store.register(password="secret-123")
    with pytest.raises(ValueError, match="invalid email"):
        store.register(email="not-an-email", password="secret-123")
    with pytest.raises(ValueError, match="at least 8"):
        store.register(email="a@b.co", password="short")


def test_register_rejects_duplicate_identifiers():
    store = AccountStore()
    store.register(email="dup@example.com", password="secret-123")
    with pytest.raises(ValueError, match="already exists"):
        store.register(email="DUP@example.com", password="other-secret")


def test_login_by_email_or_wallet():
    store = AccountStore()
    account = store.register(
        email="both@example.com", wallet="WalletXYZ", password="secret-123"
    )
    assert store.verify_login("both@example.com", "secret-123")["account_id"] == account["account_id"]
    assert store.verify_login("WalletXYZ", "secret-123")["account_id"] == account["account_id"]
    assert store.verify_login("both@example.com", "wrong-password") is None
    assert store.verify_login("nobody@example.com", "secret-123") is None


def test_sessions_resolve_and_destroy():
    store = AccountStore()
    account = store.register(email="s@example.com", password="secret-123")
    token = store.create_session(account["account_id"])
    assert store.session_account(token)["account_id"] == account["account_id"]
    store.destroy_session(token)
    assert store.session_account(token) is None
    assert store.session_account("bogus") is None


def test_api_key_lifecycle():
    store = AccountStore()
    account = store.register(email="k@example.com", password="secret-123")
    other = store.register(email="other@example.com", password="secret-123")
    key = store.create_api_key(account["account_id"])
    assert key.startswith("sk-mesh-")
    assert store.keys_for(account["account_id"]) == [key]
    # someone else's account cannot revoke it
    assert store.revoke_api_key(other["account_id"], key) is False
    assert store.revoke_api_key(account["account_id"], key) is True
    assert store.keys_for(account["account_id"]) == []
    assert store.is_key_revoked(key)


def test_accounts_persist_across_restart(tmp_path):
    db = str(tmp_path / "accounts.db")
    store = AccountStore(db_path=db)
    account = store.register(email="p@example.com", password="secret-123")
    key = store.create_api_key(account["account_id"])
    store.save_to_db()

    reloaded = AccountStore(db_path=db)
    assert reloaded.verify_login("p@example.com", "secret-123") is not None
    assert reloaded.keys_for(account["account_id"]) == [key]


def test_account_events_replicate_and_dedupe():
    leader = AccountStore()
    follower = AccountStore()
    account = leader.register(email="r@example.com", password="secret-123")
    key = leader.create_api_key(account["account_id"])
    leader.revoke_api_key(account["account_id"], key)

    events, cursor = leader.events_since(0)
    assert follower.apply_events(events) == len(events)
    assert follower.apply_events(events) == 0  # replay is a no-op
    assert follower.verify_login("r@example.com", "secret-123") is not None
    assert follower.is_key_revoked(key)
    more, _ = leader.events_since(cursor)
    assert more == []


# ---------------------------------------------------------- HTTP integration


def _call(url, method="GET", body=None, token=None):
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    with urllib.request.urlopen(req) as r:
        return json.loads(r.read())


@pytest.fixture
def account_tracker():
    ledger = BillingLedger(starting_credit=0.5, default_price_per_1k=0.02)
    tracker = TrackerServer(billing=ledger, accounts=AccountStore(), hive_secret=HIVE_SECRET)
    port = tracker.start()
    yield f"http://127.0.0.1:{port}", ledger
    tracker.stop()


def test_register_login_and_account_view(account_tracker):
    url, _ = account_tracker
    reg = _call(f"{url}/v1/auth/register", "POST",
                {"email": "admin@example.com", "password": "secret-123"})
    assert reg["account"]["role"] == "admin"
    assert reg["api_key"].startswith("sk-mesh-")
    assert reg["session_token"]

    login = _call(f"{url}/v1/auth/login", "POST",
                  {"identifier": "admin@example.com", "password": "secret-123"})
    me = _call(f"{url}/v1/account", token=login["session_token"])
    assert me["account"]["email"] == "admin@example.com"
    assert me["api_keys"] == [reg["api_key"]]
    # registration granted Caller Credit on the ledger
    assert me["total_balance"] == pytest.approx(0.5)
    assert me["usage"]["requests"] == 0


def test_bad_credentials_and_missing_session_are_401(account_tracker):
    url, _ = account_tracker
    _call(f"{url}/v1/auth/register", "POST",
          {"email": "a@example.com", "password": "secret-123"})
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _call(f"{url}/v1/auth/login", "POST",
              {"identifier": "a@example.com", "password": "wrong-pass"})
    assert exc_info.value.code == 401
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _call(f"{url}/v1/account")
    assert exc_info.value.code == 401


def test_key_create_revoke_and_revoked_key_rejected_by_proxy(account_tracker):
    url, _ = account_tracker
    reg = _call(f"{url}/v1/auth/register", "POST",
                {"email": "k@example.com", "password": "secret-123"})
    token = reg["session_token"]

    new_key = _call(f"{url}/v1/account/keys", "POST", {}, token=token)["api_key"]
    me = _call(f"{url}/v1/account", token=token)
    assert sorted(me["api_keys"]) == sorted([reg["api_key"], new_key])

    _call(f"{url}/v1/account/keys/revoke", "POST", {"api_key": new_key}, token=token)
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _call(f"{url}/v1/chat/completions", "POST",
              {"model": "any", "messages": []}, token=new_key)
    assert exc_info.value.code == 401
    assert "revoked" in exc_info.value.read().decode()


def test_admin_listing_requires_admin_role(account_tracker):
    url, _ = account_tracker
    admin = _call(f"{url}/v1/auth/register", "POST",
                  {"email": "admin@example.com", "password": "secret-123"})
    user = _call(f"{url}/v1/auth/register", "POST",
                 {"wallet": "WalletUser1", "password": "secret-123"})

    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _call(f"{url}/v1/admin/accounts", token=user["session_token"])
    assert exc_info.value.code == 403

    listing = _call(f"{url}/v1/admin/accounts", token=admin["session_token"])
    accounts = listing["accounts"]
    assert len(accounts) == 2
    assert accounts[0]["role"] == "admin"
    assert accounts[1]["wallet"] == "WalletUser1"
    assert "balances" in accounts[0]


def test_accounts_gossip_endpoint_applies_events(account_tracker):
    url, _ = account_tracker
    peer = AccountStore()
    peer.register(email="remote@example.com", password="secret-123")
    events, _ = peer.events_since(0)
    body = json.dumps({"events": events}).encode()
    req = urllib.request.Request(
        f"{url}/v1/accounts/gossip", data=body,
        headers={"Content-Type": "application/json", **sign_hive_request(HIVE_SECRET, body)},
        method="POST",
    )
    with urllib.request.urlopen(req) as r:
        result = json.loads(r.read())
    assert result["applied"] == len(events)
    login = _call(f"{url}/v1/auth/login", "POST",
                  {"identifier": "remote@example.com", "password": "secret-123"})
    assert login["account"]["email"] == "remote@example.com"


def test_accounts_endpoints_404_when_disabled():
    tracker = TrackerServer()  # no accounts, no billing
    port = tracker.start()
    try:
        with pytest.raises(urllib.error.HTTPError) as exc_info:
            _call(f"http://127.0.0.1:{port}/v1/auth/register", "POST",
                  {"email": "x@example.com", "password": "secret-123"})
        assert exc_info.value.code == 404
    finally:
        tracker.stop()