from types import SimpleNamespace

from gpu_rent.provision import _OLLAMA_INSTALL_ENV, _remote_llm_env


def test_remote_llm_env_forwards_ollama_vars(monkeypatch):
    monkeypatch.setenv("OLLAMA_VERSION", "0.6.5")
    monkeypatch.setenv("OLLAMA_SHA256", "abc123")
    cfg = SimpleNamespace(ssh_user="ubuntu")
    env = _remote_llm_env(cfg, *_OLLAMA_INSTALL_ENV)
    assert env["SWARM_USER"] == "ubuntu"
    assert env["OLLAMA_VERSION"] == "0.6.5"
    assert env["OLLAMA_SHA256"] == "abc123"


def test_remote_llm_env_skips_empty(monkeypatch):
    for key in _OLLAMA_INSTALL_ENV:
        monkeypatch.delenv(key, raising=False)
    cfg = SimpleNamespace(ssh_user="ubuntu")
    env = _remote_llm_env(cfg, *_OLLAMA_INSTALL_ENV)
    assert env == {"SWARM_USER": "ubuntu"}