TradeBot/tests/test_training_coordination.py

from __future__ import annotations

import base64
import hashlib
import json

import pytest

from crypto_spot_bot.training_coordination import TrainingCoordinator


def test_training_coordinator_claims_and_completes_job(tmp_path) -> None:
    coordinator = TrainingCoordinator(tmp_path)

    requested = coordinator.request_retrain({"source": "android"})
    job_id = requested["job"]["id"]
    heartbeat = coordinator.heartbeat({"worker_id": "win-1", "name": "DESKTOP-TMFDL0H"})
    claimed = coordinator.claim({"worker_id": "win-1", "name": "DESKTOP-TMFDL0H"})

    assert requested["queued"] is True
    assert heartbeat["status"]["agent_online"] is True
    assert claimed["claimed"] is True
    assert claimed["job"]["id"] == job_id
    assert coordinator.status()["active_job"]["status"] == "running"

    progress = coordinator.progress(
        job_id,
        {"status": "running", "phase": "training", "progress_percent": 42, "message": "epoch 1"},
    )

    assert progress["job"]["phase"] == "training"
    assert progress["job"]["progress_percent"] == 42
    assert coordinator.status()["active_job"]["message"] == "epoch 1"

    completed = coordinator.complete(job_id, {"success": True, "message": "ok"})

    assert completed["job"]["status"] == "completed"
    assert coordinator.status()["active_job"] is None


def test_training_coordinator_preserves_boolean_resume_candidate_parameter(tmp_path) -> None:
    coordinator = TrainingCoordinator(tmp_path)

    requested = coordinator.request_retrain(
        {"source": "recovery", "parameters": {"resume_candidate": True}}
    )

    assert requested["job"]["parameters"] == {"resume_candidate": True}


def test_training_coordinator_accepts_chunked_artifact_upload(tmp_path) -> None:
    coordinator = TrainingCoordinator(tmp_path)
    job = coordinator.request_retrain({"source": "test"})["job"]
    coordinator.claim({"worker_id": "test-worker"})
    payload = b'{"type":"pytorch_recurrent_forecaster","symbols":{}}\n'
    sha256 = hashlib.sha256(payload).hexdigest()
    first = payload[:20]
    second = payload[20:]

    part_1 = coordinator.save_artifact_chunk(
        job["id"],
        {
            "name": "lstm_forecaster.json",
            "index": 0,
            "total": 2,
            "sha256": sha256,
            "data_base64": base64.b64encode(first).decode("ascii"),
        },
    )
    part_2 = coordinator.save_artifact_chunk(
        job["id"],
        {
            "name": "lstm_forecaster.json",
            "index": 1,
            "total": 2,
            "sha256": sha256,
            "data_base64": base64.b64encode(second).decode("ascii"),
        },
    )

    assert part_1["complete"] is False
    assert part_2["complete"] is True
    assert not (tmp_path / "lstm_forecaster.json").exists()
    assert (tmp_path / ".training_uploads" / job["id"] / "ready" / "lstm_forecaster.json").read_bytes() == payload
    assert coordinator.status()["latest_job"]["artifacts"][0]["sha256"] == sha256


def test_running_claimed_job_keeps_agent_online_when_heartbeat_is_stale(tmp_path) -> None:
    coordinator = TrainingCoordinator(tmp_path)
    coordinator.request_retrain({"source": "android"})
    coordinator.claim({"worker_id": "win-1", "name": "DESKTOP-TMFDL0H"})

    state_path = tmp_path / "training_coordination.json"
    state = json.loads(state_path.read_text(encoding="utf-8"))
    state["worker"]["last_seen_at"] = "2026-01-01T00:00:00+00:00"
    state_path.write_text(json.dumps(state), encoding="utf-8")

    status = coordinator.status()

    assert status["agent_recently_seen"] is False
    assert status["agent_busy"] is True
    assert status["agent_online"] is True


def test_training_upload_rejects_unknown_job(tmp_path) -> None:
    coordinator = TrainingCoordinator(tmp_path)
    payload = b"{}"

    with pytest.raises(ValueError, match="not found"):
        coordinator.save_artifact_chunk(
            "11111111-1111-4111-8111-111111111111",
            {
                "name": "lstm_forecaster.json",
                "index": 0,
                "total": 1,
                "sha256": hashlib.sha256(payload).hexdigest(),
                "data_base64": base64.b64encode(payload).decode("ascii"),
            },
        )


def test_training_bundle_promotes_only_after_successful_guard(tmp_path) -> None:
    coordinator = TrainingCoordinator(tmp_path)
    job = coordinator.request_retrain({"source": "test"})["job"]
    coordinator.claim({"worker_id": "worker-1"})
    model = {
        "type": "pytorch_recurrent_forecaster",
        "symbols": {
            "BTCUSDT": {
                "model": "torch_gru",
                "lookback": 4,
                "input_size": 1,
                "hidden_size": 1,
                "state_dict": {"weight_ih_l0": [[0.0]]},
                "head_weight": [[0.0]],
                "head_bias": [0.0],
            }
        },
    }
    model_payload = (json.dumps(model) + "\n").encode()
    model_sha256 = hashlib.sha256(model_payload).hexdigest()
    artifacts = {
        "lstm_forecaster.json": model,
        "torch_retrain_guard.json": {
            "accepted": True,
            "candidate_artifact_sha256": model_sha256,
        },
        "torch_threshold_calibration.json": {
            "artifact_sha256": model_sha256,
            "validation": {
                "passed": True,
                "protocol": "untouched_model_holdout_with_threshold_walk_forward",
            }
        },
    }
    for name, data in artifacts.items():
        payload = model_payload if name == "lstm_forecaster.json" else (json.dumps(data) + "\n").encode()
        coordinator.save_artifact_chunk(
            job["id"],
            {
                "name": name,
                "index": 0,
                "total": 1,
                "sha256": hashlib.sha256(payload).hexdigest(),
                "data_base64": base64.b64encode(payload).decode("ascii"),
            },
        )

    completed = coordinator.complete(job["id"], {"success": True})

    assert completed["job"]["status"] == "completed"
    assert json.loads((tmp_path / "lstm_forecaster.json").read_text())["symbols"]["BTCUSDT"]