SciMesh/tests/test_worker_auth.py

"""Tests for worker token strategies and the client's 401 refresh."""

from __future__ import annotations

import json
from pathlib import Path
from urllib.error import HTTPError

import pytest

from scimesh.worker.auth import (
    StaticTokenProvider,
    TokenExchangeError,
    WorkerKeyTokenProvider,
    provider_from_config,
)
from scimesh.worker.config import WorkerConfig
from scimesh.worker.coordinator import HttpCoordinatorClient


class FakeResponse:
    def __init__(self, status: int, body: bytes) -> None:
        self.status = status
        self._body = body

    def read(self) -> bytes:
        return self._body

    def __enter__(self) -> "FakeResponse":
        return self

    def __exit__(self, *exc) -> bool:
        return False


class SeqOpener:
    """Returns/raises a scripted sequence of responses, recording each request."""

    def __init__(self, actions: list) -> None:
        self.actions = list(actions)
        self.requests: list = []

    def open(self, request, timeout=None):
        self.requests.append(request)
        action = self.actions.pop(0)
        if isinstance(action, Exception):
            raise action
        return action


def _exchange_response(token: str, expires_in: int) -> FakeResponse:
    return FakeResponse(200, json.dumps({"token": token, "expires_in": expires_in}).encode())


def test_static_provider_returns_fixed_token_and_never_refreshes():
    provider = StaticTokenProvider("tok")
    assert provider.token() == "tok"
    provider.refresh()
    assert provider.token() == "tok"


def test_static_provider_none_means_no_auth():
    assert StaticTokenProvider(None).token() is None


def test_worker_key_provider_exchanges_once_then_caches():
    clock = {"t": 1000.0}
    provider = WorkerKeyTokenProvider(
        "http://users", "scimesh_wk_live_x", timeout=5, now=lambda: clock["t"]
    )
    provider._opener = SeqOpener([_exchange_response("jwt-1", 100)])

    # First call exchanges; a second call well within the TTL reuses the cache.
    assert provider.token() == "jwt-1"
    clock["t"] = 1050.0  # 50s later, TTL 100s with 0.2 leeway → refresh at +80s
    assert provider.token() == "jwt-1"
    assert len(provider._opener.requests) == 1


def test_worker_key_provider_refreshes_after_leeway():
    clock = {"t": 0.0}
    provider = WorkerKeyTokenProvider(
        "http://users", "k", timeout=5, now=lambda: clock["t"]
    )
    provider._opener = SeqOpener([
        _exchange_response("jwt-1", 100),
        _exchange_response("jwt-2", 100),
    ])
    assert provider.token() == "jwt-1"
    clock["t"] = 85.0  # past the 80s refresh point
    assert provider.token() == "jwt-2"
    assert len(provider._opener.requests) == 2


def test_worker_key_provider_force_refresh():
    provider = WorkerKeyTokenProvider("http://users", "k", timeout=5, now=lambda: 0.0)
    provider._opener = SeqOpener([
        _exchange_response("jwt-1", 100),
        _exchange_response("jwt-2", 100),
    ])
    assert provider.token() == "jwt-1"
    provider.refresh()
    assert provider.token() == "jwt-2"


def test_worker_key_provider_raises_on_rejected_key():
    provider = WorkerKeyTokenProvider("http://users", "bad", timeout=5, now=lambda: 0.0)
    provider._opener = SeqOpener([HTTPError("http://users", 401, "unauthorized", {}, None)])
    with pytest.raises(TokenExchangeError):
        provider.token()


def test_worker_key_provider_raises_when_token_missing():
    provider = WorkerKeyTokenProvider("http://users", "k", timeout=5, now=lambda: 0.0)
    provider._opener = SeqOpener([FakeResponse(200, json.dumps({"expires_in": 100}).encode())])
    with pytest.raises(TokenExchangeError):
        provider.token()


def test_provider_from_config_selects_worker_key_mode():
    provider = provider_from_config(
        worker_key="scimesh_wk_live_x",
        userservice_url="http://users",
        bearer_token="ignored",
        request_timeout=5,
    )
    assert isinstance(provider, WorkerKeyTokenProvider)


def test_provider_from_config_falls_back_to_static():
    provider = provider_from_config(
        worker_key=None, userservice_url=None, bearer_token="tok", request_timeout=5
    )
    assert isinstance(provider, StaticTokenProvider)
    assert provider.token() == "tok"


class RefreshCountingProvider:
    def __init__(self) -> None:
        self.tokens = ["stale", "fresh"]
        self.index = 0
        self.refreshes = 0

    def token(self) -> str:
        return self.tokens[min(self.index, len(self.tokens) - 1)]

    def refresh(self) -> None:
        self.refreshes += 1
        self.index += 1


def test_coordinator_client_refreshes_and_retries_once_on_401():
    provider = RefreshCountingProvider()
    client = HttpCoordinatorClient("http://coord", timeout=5, token_provider=provider)
    client._opener = SeqOpener([
        HTTPError("http://coord/tasks/claim", 401, "unauthorized", {}, None),
        FakeResponse(204, b""),
    ])

    status, _ = client._request("POST", "/tasks/claim", {"worker_id": "w"})

    assert status == 204
    assert provider.refreshes == 1
    # The retry carried the refreshed token.
    assert provider.index == 1


def test_coordinator_client_does_not_loop_on_persistent_401():
    provider = RefreshCountingProvider()
    client = HttpCoordinatorClient("http://coord", timeout=5, token_provider=provider)
    client._opener = SeqOpener([
        HTTPError("http://coord/x", 401, "unauthorized", {}, None),
        HTTPError("http://coord/x", 401, "unauthorized", {}, None),
    ])

    status, _ = client._request("POST", "/x", {})

    # One refresh, one retry, then the second 401 is surfaced rather than retried.
    assert status == 401
    assert provider.refreshes == 1


def _base_config(**extra) -> dict:
    return {
        "coordinator_url": "http://coord",
        "worker_id": None,
        "work_dir": Path("."),
        **extra,
    }


def test_worker_key_requires_userservice_url():
    with pytest.raises(ValueError, match="userservice_url"):
        WorkerConfig(**_base_config(worker_key="scimesh_wk_live_x"))


def test_worker_key_with_userservice_url_is_valid():
    cfg = WorkerConfig(**_base_config(worker_key="scimesh_wk_live_x", userservice_url="http://users"))
    assert cfg.worker_key == "scimesh_wk_live_x"
    assert cfg.userservice_url == "http://users"


def test_userservice_url_must_be_absolute():
    with pytest.raises(ValueError, match="userservice_url"):
        WorkerConfig(**_base_config(userservice_url="not-a-url"))