neuron-tai/tests/test_mining_cli.py

"""Tests for US-016: mining-style node startup CLI + live dashboard."""

from __future__ import annotations

import json
import socket
import sys
import types
from pathlib import Path
from unittest.mock import MagicMock, patch

# A fake node server has no real backend to prove capability with; say so
# explicitly rather than bypassing startup's fail-closed admission.
from meshnet_node.testing import assume_capability


# ---------------------------------------------------------------------------
# model_catalog tests
# ---------------------------------------------------------------------------

def test_curated_models_list_is_non_empty():
    "Curated models list is non empty\n\nTags: general"
    from meshnet_node.model_catalog import CURATED_MODELS
    assert len(CURATED_MODELS) >= 5


def test_model_preset_vram_for_quant():
    "Model preset vram for quant\n\nTags: general"
    from meshnet_node.model_catalog import CURATED_MODELS

    m = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name)
    assert m.vram_for_quant("nf4") == m.vram_nf4
    assert m.vram_for_quant("int8") == m.vram_int8
    assert m.vram_for_quant("bf16") == m.vram_bf16
    assert m.vram_for_quant("bfloat16") == m.vram_bf16  # alias


def test_model_preset_fits_vram():
    "Model preset fits vram\n\nTags: general"
    from meshnet_node.model_catalog import CURATED_MODELS

    small = next(m for m in CURATED_MODELS if m.vram_nf4 < 10)
    assert small.fits_vram(small.vram_nf4, "nf4")
    assert not small.fits_vram(small.vram_nf4 - 1, "nf4")


def test_recommended_quant_respects_vram():
    "Recommended quant respects vram\n\nTags: general"
    from meshnet_node.model_catalog import CURATED_MODELS

    m = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name)
    # nf4=18, int8=40, bf16=140
    assert m.recommended_quant(200) == "bf16"
    assert m.recommended_quant(50) == "int8"
    assert m.recommended_quant(20) == "nf4"
    assert m.recommended_quant(5) is None


def test_models_with_insufficient_vram_are_marked(monkeypatch):
    "Models with insufficient vram are marked\n\nTags: general"
    from meshnet_node import wizard as wiz

    # Simulate 6 GB GPU
    gpus = [{"index": 0, "name": "RTX 3060", "vram_gb": 6.0, "backend": "cuda"}]
    monkeypatch.setattr(wiz, "_detect_gpus", lambda: gpus)

    # Phi-3 at NF4 needs 4 GB — should fit; Llama-3-70B at NF4 needs 18 GB — should not
    from meshnet_node.model_catalog import CURATED_MODELS

    phi = next(m for m in CURATED_MODELS if "Phi-3" in m.name)
    llama = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name)

    assert phi.fits_vram(6.0, "nf4")
    assert not llama.fits_vram(6.0, "nf4")


# ---------------------------------------------------------------------------
# config tests
# ---------------------------------------------------------------------------

def test_load_config_returns_none_when_missing(tmp_path):
    "Load config returns none when missing\n\nTags: general"
    from meshnet_node.config import load_config
    assert load_config(tmp_path / "nonexistent.json") is None


def test_save_and_load_config_roundtrip(tmp_path):
    "Save and load config roundtrip\n\nTags: general"
    from meshnet_node.config import save_config, load_config

    cfg = {"model_hf_repo": "test/model", "quantization": "nf4", "tracker_url": "http://localhost:8080"}
    cfg_path = tmp_path / "config.json"
    save_config(cfg, cfg_path)

    loaded = load_config(cfg_path)
    assert loaded == cfg


def test_save_config_creates_parent_dirs(tmp_path):
    "Save config creates parent dirs\n\nTags: general"
    from meshnet_node.config import save_config, load_config

    nested = tmp_path / "deep" / "nested" / "config.json"
    save_config({"x": 1}, nested)
    assert nested.exists()
    assert load_config(nested) == {"x": 1}


def test_merge_cli_overrides_applies_non_none_values():
    "Merge cli overrides applies non none values\n\nTags: general"
    from meshnet_node.config import merge_cli_overrides

    base = {"tracker_url": "http://a:8080", "quantization": "nf4", "port": 7000}
    result = merge_cli_overrides(base, tracker_url="http://b:9090", port=None)
    assert result["tracker_url"] == "http://b:9090"
    assert result["port"] == 7000  # None override ignored
    assert result["quantization"] == "nf4"  # unchanged


# ---------------------------------------------------------------------------
# wizard tests
# ---------------------------------------------------------------------------

def test_print_models_table_runs_without_error(capsys, monkeypatch):
    "Print models table runs without error\n\nTags: general"
    from meshnet_node import wizard as wiz

    monkeypatch.setattr(wiz, "_detect_gpus", lambda: [{"index": 0, "name": "GPU", "vram_gb": 24.0, "backend": "cuda"}])
    wiz.print_models_table()
    out = capsys.readouterr().out
    assert "Llama" in out or "Qwen" in out or "Phi" in out


def test_wizard_writes_config_on_happy_path(tmp_path, monkeypatch):
    "Wizard writes config on happy path\n\nTags: general"
    from meshnet_node import wizard as wiz
    from meshnet_node.config import load_config, save_config

    # Fake GPU
    gpus = [{"index": 0, "name": "RTX 4090", "vram_gb": 24.0, "backend": "cuda"}]
    monkeypatch.setattr(wiz, "_detect_gpus", lambda: gpus)
    # Tracker not reachable (stub)
    monkeypatch.setattr(wiz, "_ping_tracker", lambda url: False)

    # Simulate user selecting model 1 (Qwen2.5-0.5B), quant 1 (nf4), default dir, default tracker, default wallet
    inputs = iter([
        "1",          # pick Qwen2.5-0.5B-Instruct (index 1 in CURATED_MODELS)
        "1",          # quant NF4
        str(tmp_path / "models"),  # download dir
        "http://localhost:8080",   # tracker
        str(tmp_path / "wallet.json"),  # wallet
    ])
    monkeypatch.setattr("builtins.input", lambda prompt="": next(inputs))

    cfg = wiz.run_wizard(config_path_override=tmp_path / "config.json")
    assert cfg["model_hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct"
    assert cfg["quantization"] == "nf4"
    assert "download_dir" in cfg
    assert cfg["tracker_url"] == "http://localhost:8080"


def test_wizard_raises_keyboard_interrupt_on_ctrl_c(monkeypatch):
    "Wizard raises keyboard interrupt on ctrl c\n\nTags: general"
    from meshnet_node import wizard as wiz

    gpus = [{"index": 0, "name": "RTX 4090", "vram_gb": 24.0, "backend": "cuda"}]
    monkeypatch.setattr(wiz, "_detect_gpus", lambda: gpus)

    call_count = [0]

    def fake_input(prompt=""):
        call_count[0] += 1
        if call_count[0] == 1:
            raise KeyboardInterrupt

    monkeypatch.setattr("builtins.input", fake_input)

    import pytest
    with pytest.raises(KeyboardInterrupt):
        wiz.run_wizard()


# ---------------------------------------------------------------------------
# dashboard tests
# ---------------------------------------------------------------------------

def test_is_interactive_tty_false_when_not_tty(monkeypatch):
    "Is interactive tty false when not tty\n\nTags: general"
    from meshnet_node import dashboard as dash

    monkeypatch.setattr(sys.stdout, "isatty", lambda: False)
    assert not dash.is_interactive_tty()


def test_dashboard_plain_fallback_on_keyboard_interrupt(monkeypatch):
    "Plain loop exits cleanly when Ctrl-C is raised.\n\nTags: general"
    from meshnet_node import dashboard as dash

    node = MagicMock()
    node.chat_completion_count = 5

    call_count = [0]

    def fake_sleep(t):
        call_count[0] += 1
        if call_count[0] >= 1:
            raise KeyboardInterrupt

    monkeypatch.setattr(dash.time, "sleep", fake_sleep)
    monkeypatch.setattr(dash, "_gpu_stats", lambda: [])
    monkeypatch.setattr(sys.stdout, "isatty", lambda: False)

    cfg = {"model_name": "test-model", "quantization": "nf4"}
    # Should not raise
    dash.run_dashboard(node, cfg, start_time=dash.time.monotonic())


def test_ema_updates_correctly():
    "Ema updates correctly\n\nTags: general"
    from meshnet_node.dashboard import _EMA

    ema = _EMA(alpha=1.0)  # alpha=1.0 → always takes latest sample
    ema.update(10.0)
    assert ema.value == 10.0
    ema.update(20.0)
    assert ema.value == 20.0


# ---------------------------------------------------------------------------
# CLI integration tests
# ---------------------------------------------------------------------------

def test_models_command_prints_table(capsys, monkeypatch):
    "meshnet-node models prints the curated table and exits 0.\n\nTags: general"
    from meshnet_node import wizard as wiz

    monkeypatch.setattr(wiz, "_detect_gpus", lambda: [])

    from meshnet_node.cli import main
    monkeypatch.setattr(sys, "argv", ["meshnet-node", "models"])

    try:
        main()
    except SystemExit as exc:
        assert exc.code == 0

    out = capsys.readouterr().out
    assert "Llama" in out or "Qwen" in out or "Phi" in out


def test_config_command_no_config_exits_1(tmp_path, monkeypatch):
    "Config command no config exits 1\n\nTags: general"
    from meshnet_node import config as cfg_mod
    from meshnet_node.cli import main

    monkeypatch.setattr(cfg_mod, "_DEFAULT_CONFIG_FILE", tmp_path / "nonexistent.json")
    monkeypatch.setattr(sys, "argv", ["meshnet-node", "config"])

    with patch("meshnet_node.config.config_path", return_value=tmp_path / "nonexistent.json"):
        try:
            main()
        except SystemExit as exc:
            assert exc.code == 1


def test_config_command_prints_saved_config(tmp_path, monkeypatch, capsys):
    "Config command prints saved config\n\nTags: general"
    from meshnet_node import config as cfg_mod
    from meshnet_node.config import save_config
    from meshnet_node.cli import main

    saved = {"model_hf_repo": "meta-llama/Meta-Llama-3-70B-Instruct", "quantization": "nf4"}
    cfg_file = tmp_path / "config.json"
    save_config(saved, cfg_file)

    monkeypatch.setattr(sys, "argv", ["meshnet-node", "config"])

    with patch("meshnet_node.config.config_path", return_value=cfg_file):
        with patch("meshnet_node.config.load_config", return_value=saved):
            try:
                main()
            except SystemExit as exc:
                assert exc.code == 0

    out = capsys.readouterr().out
    data = json.loads(out.split("\n", 1)[1])  # skip the "Config: ..." header line
    assert data["model_hf_repo"] == saved["model_hf_repo"]


def test_detect_num_layers_returns_catalog_value_without_network(monkeypatch):
    "detect_num_layers uses the curated catalog first — no network call.\n\nTags: general"
    from meshnet_node.model_catalog import detect_num_layers

    # Qwen2.5-0.5B is in the catalog with 24 layers
    layers = detect_num_layers("Qwen/Qwen2.5-0.5B-Instruct")
    assert layers == 24


def test_detect_num_layers_returns_none_on_error(monkeypatch):
    "Detect num layers returns none on error\n\nTags: general"
    from meshnet_node.model_catalog import detect_num_layers

    # Monkeypatch AutoConfig to raise
    import meshnet_node.model_catalog as cat
    monkeypatch.setattr(cat, "detect_num_layers", lambda repo: None if "bad" in repo else detect_num_layers(repo))
    assert cat.detect_num_layers("bad/repo") is None


def test_startup_auto_detects_shard_range(monkeypatch, tmp_path):
    "When shard_start/end are None, startup reads layer count from catalog.\n\nTags: general"
    from meshnet_node import startup as su
    from meshnet_node.model_catalog import detect_num_layers

    calls = []

    def fake_detect(repo):
        calls.append(repo)
        return 24  # Qwen2.5-0.5B

    monkeypatch.setattr(su, "_detect_num_layers", fake_detect)

    # Fake hardware detection
    monkeypatch.setattr(su, "detect_hardware", lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0})

    # Fake wallet
    monkeypatch.setattr(su, "load_or_create_wallet", lambda **kw: (None, None, "fake-wallet"))

    # Fake TorchNodeServer
    class FakeNode:
        chat_completion_count = 0
        def start(self): return 9999
        def stop(self): pass

    import meshnet_node.startup as su2
    monkeypatch.setattr(su2, "TorchNodeServer", lambda **kw: FakeNode())

    node = su.run_startup(
        tracker_url="http://localhost:8080",
        model_id="Qwen/Qwen2.5-0.5B-Instruct",
        # shard_start and shard_end intentionally omitted
        quantization="bfloat16",
        host="127.0.0.1",
        capability_validator=assume_capability,
    )
    assert calls == ["Qwen/Qwen2.5-0.5B-Instruct"]
    assert isinstance(node, FakeNode)


def test_legacy_start_subcommand_accepted(monkeypatch):
    "meshnet-node start --tracker http://...\n\nTags: general"
    from meshnet_node.cli import main

    def fake_run_startup(*args, **kwargs):
        class _FakeNode:
            chat_completion_count = 0
            def stop(self): pass
        return _FakeNode()

    monkeypatch.setattr(sys, "argv", [
        "meshnet-node", "start",
        "--tracker", "http://localhost:8080",
        "--model", "stub-model",
        "--port", "0",
    ])

    raised = []

    def fake_sleep(t):
        raise KeyboardInterrupt

    with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
        with patch("time.sleep", side_effect=fake_sleep):
            try:
                main()
            except SystemExit as exc:
                raised.append(exc.code)

    # Exited (either 0 or via KeyboardInterrupt caught in _cmd_start)
    # The important thing is no unhandled exception from arg parsing


def test_legacy_start_treats_repo_model_as_model_id(monkeypatch):
    "`meshnet-node start --model org/repo` enters the HF model startup path.\n\nTags: general"
    from meshnet_node.cli import main

    captured = {}

    def fake_run_startup(*args, **kwargs):
        captured.update(kwargs)
        class _FakeNode:
            chat_completion_count = 0
            def stop(self): pass
        return _FakeNode()

    monkeypatch.setattr(sys, "argv", [
        "meshnet-node", "start",
        "--tracker", "http://192.168.0.179:8081",
        "--model", "Qwen/Qwen2.5-0.5B-Instruct",
        "--port", "0",
    ])

    with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
        with patch("time.sleep", side_effect=KeyboardInterrupt):
            try:
                main()
            except SystemExit as exc:
                assert exc.code == 0

    assert captured["model"] == "Qwen2.5-0.5B-Instruct"
    assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"


def test_legacy_start_catalog_model_with_pinned_shards(monkeypatch):
    "Catalog model names accept --shard-start/--shard-end without --model-id.\n\nTags: general"
    from meshnet_node.cli import main

    captured = {}

    def fake_run_startup(*args, **kwargs):
        captured.update(kwargs)
        class _FakeNode:
            chat_completion_count = 0
            def stop(self): pass
        return _FakeNode()

    monkeypatch.setattr(sys, "argv", [
        "meshnet-node", "start",
        "--tracker", "http://192.168.0.179:8080",
        "--model", "Qwen3.6-35B-A3B",
        "--shard-start", "0",
        "--shard-end", "44",
        "--port", "0",
    ])

    with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
        with patch("time.sleep", side_effect=KeyboardInterrupt):
            try:
                main()
            except SystemExit as exc:
                assert exc.code == 0

    assert captured["model"] == "Qwen3.6-35B-A3B"
    assert captured["model_id"] is None
    assert captured["shard_start"] == 0
    assert captured["shard_end"] == 44


def test_legacy_start_model_id_alias_for_catalog_name(monkeypatch):
    "--model-id with a catalog name routes through the tracker preset path.\n\nTags: general"
    from meshnet_node.cli import main

    captured = {}

    def fake_run_startup(*args, **kwargs):
        captured.update(kwargs)
        class _FakeNode:
            chat_completion_count = 0
            def stop(self): pass
        return _FakeNode()

    monkeypatch.setattr(sys, "argv", [
        "meshnet-node", "start",
        "--tracker", "http://192.168.0.179:8080",
        "--model-id", "Qwen3.6-35B-A3B",
        "--port", "0",
    ])

    with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
        with patch("time.sleep", side_effect=KeyboardInterrupt):
            try:
                main()
            except SystemExit as exc:
                assert exc.code == 0

    assert captured["model"] == "Qwen3.6-35B-A3B"
    assert captured["model_id"] is None


def test_legacy_start_hf_repo_with_pinned_shards(monkeypatch):
    "HF repo --model with pinned shards still enters the torch startup path.\n\nTags: general"
    from meshnet_node.cli import main

    captured = {}

    def fake_run_startup(*args, **kwargs):
        captured.update(kwargs)
        class _FakeNode:
            chat_completion_count = 0
            def stop(self): pass
        return _FakeNode()

    monkeypatch.setattr(sys, "argv", [
        "meshnet-node", "start",
        "--tracker", "http://192.168.0.179:8081",
        "--model", "Qwen/Qwen2.5-0.5B-Instruct",
        "--shard-start", "12",
        "--shard-end", "23",
        "--port", "0",
    ])

    with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
        with patch("time.sleep", side_effect=KeyboardInterrupt):
            try:
                main()
            except SystemExit as exc:
                assert exc.code == 0

    assert captured["model"] == "Qwen2.5-0.5B-Instruct"
    assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
    assert captured["shard_start"] == 12
    assert captured["shard_end"] == 23


def test_legacy_start_falls_back_to_env_tracker_and_model(monkeypatch):
    "`meshnet-node start` uses env defaults when tracker/model flags are omitted.\n\nTags: general"
    import importlib

    from meshnet_node import config as config_mod
    from meshnet_node.cli import main

    monkeypatch.setenv("MESHNET_TRACKER_URL", "http://env-tracker:8081")
    monkeypatch.setenv("MESHNET_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
    importlib.reload(config_mod)

    captured = {}

    def fake_run_startup(*args, **kwargs):
        captured.update(kwargs)
        class _FakeNode:
            chat_completion_count = 0
            def stop(self): pass
        return _FakeNode()

    monkeypatch.setattr(sys, "argv", [
        "meshnet-node", "start",
        "--port", "0",
    ])

    try:
        with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
            with patch("time.sleep", side_effect=KeyboardInterrupt):
                try:
                    main()
                except SystemExit as exc:
                    assert exc.code == 0
    finally:
        monkeypatch.delenv("MESHNET_TRACKER_URL", raising=False)
        monkeypatch.delenv("MESHNET_MODEL", raising=False)
        importlib.reload(config_mod)

    assert captured["tracker_url"] == "http://env-tracker:8081"
    assert captured["model"] == "Qwen2.5-0.5B-Instruct"
    assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"


def test_first_available_port_skips_an_occupied_custom_port():
    "Port search skips an occupied custom base port.\n\nTags: general"
    from meshnet_node.cli import _first_available_port

    occupied = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    occupied.bind(("127.0.0.1", 0))
    start = occupied.getsockname()[1]
    occupied.listen(1)
    try:
        assert _first_available_port("127.0.0.1", start=start, attempts=100) > start
    finally:
        occupied.close()


def test_legacy_start_without_port_uses_port_search(monkeypatch):
    "Omitting --port delegates legacy startup to the port-search helper.\n\nTags: general"
    from meshnet_node import cli as cli_mod

    captured = {}

    def fake_run_startup(*args, **kwargs):
        captured.update(kwargs)
        class _FakeNode:
            chat_completion_count = 0
            def stop(self): pass
        return _FakeNode()

    monkeypatch.setattr(cli_mod, "_first_available_port", lambda *_args, **_kwargs: 7001)
    monkeypatch.setattr(sys, "argv", [
        "meshnet-node", "start",
        "--tracker", "http://192.168.0.179:8081",
        "--model", "Qwen/Qwen2.5-0.5B-Instruct",
        "--host", "127.0.0.1",
    ])

    with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
        with patch("time.sleep", side_effect=KeyboardInterrupt):
            try:
                cli_mod.main()
            except SystemExit as exc:
                assert exc.code == 0

    assert captured["port"] == 7001


def test_default_cli_passes_advertise_host(monkeypatch):
    "The documented no-subcommand LAN flag reaches startup.\n\nTags: general"
    from meshnet_node.cli import main

    saved = {
        "model_hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
        "model_name": "Qwen2.5-0.5B-Instruct",
        "quantization": "nf4",
        "tracker_url": "http://localhost:8080",
        "wallet_path": "",
        "download_dir": "",
        "port": 7000,
        "host": "0.0.0.0",
    }
    captured = {}

    def fake_run_startup(*args, **kwargs):
        captured.update(kwargs)
        class _FakeNode:
            chat_completion_count = 0
            def stop(self): pass
        return _FakeNode()

    monkeypatch.setattr(sys, "argv", [
        "meshnet-node",
        "--tracker", "http://192.168.0.179:8081",
        "--advertise-host", "192.168.0.42",
        "--debug",
        "--no-tui",
    ])

    with patch("meshnet_node.config.load_config", return_value=saved):
        with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
            with patch("meshnet_node.dashboard.run_dashboard", side_effect=KeyboardInterrupt):
                try:
                    main()
                except SystemExit as exc:
                    assert exc.code == 0

    assert captured["tracker_url"] == "http://192.168.0.179:8081"
    assert captured["advertise_host"] == "192.168.0.42"
    assert captured["debug"] is True


def test_default_cli_passes_force_cpu(monkeypatch):
    "`meshnet-node --cpu` forwards force_cpu into run_startup.\n\nTags: general"
    from meshnet_node.cli import main

    captured = {}

    def fake_run_startup(*args, **kwargs):
        captured.update(kwargs)

        class _FakeNode:
            chat_completion_count = 0

            def stop(self):
                pass

        return _FakeNode()

    saved = {
        "tracker_url": "http://localhost:8080",
        "model_name": "stub-model",
        "model_hf_repo": "",
        "quantization": "auto",
        "download_dir": "/tmp/models",
        "wallet_path": "/tmp/wallet.json",
        "port": 7000,
        "host": "0.0.0.0",
    }

    monkeypatch.setattr(sys, "argv", ["meshnet-node", "--cpu"])

    with patch("meshnet_node.config.load_config", return_value=saved):
        with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
            with patch("meshnet_node.dashboard.run_dashboard", side_effect=KeyboardInterrupt):
                try:
                    main()
                except SystemExit as exc:
                    assert exc.code == 0

    assert captured["force_cpu"] is True