"""Unit tests for assistent persona-pack seed helpers.""" from __future__ import annotations import json from pathlib import Path from unittest.mock import MagicMock from gpu_rent.provision import _safe_persona_id, seed_assistent_personas def test_safe_persona_id(): assert _safe_persona_id("neutral") == "neutral" assert _safe_persona_id("dreamer_v2") == "dreamer_v2" assert _safe_persona_id("../x") is None assert _safe_persona_id("a/b") is None assert _safe_persona_id("") is None def test_seed_assistent_personas_pushes_packs(tmp_path: Path, monkeypatch): root = tmp_path / "assistent-extensions" pack = root / "sample-pack" pack.mkdir(parents=True) (pack / "assistent-pack.yaml").write_text("kind: persona\nid: sample\n", encoding="utf-8") (pack / "persona.json").write_text( json.dumps({"title": "Sample", "accent": "#7aa2f7"}), encoding="utf-8" ) (pack / "voice.json").write_text(json.dumps({"tone": ["calm"]}), encoding="utf-8") puts: dict[str, str] = {} pushed: list[tuple[str, str]] = [] ssh_cmds: list[str] = [] def fake_put(cfg, host, remote, text, mode=0o644): puts[remote] = text def fake_ssh(cfg, host, cmd, check=False, timeout=60): ssh_cmds.append(cmd) if ".gpu-rent-gpu.json" in cmd: return json.dumps({"vram_mib": 24576, "name": "4090", "compute_cap": "8.9"}) return "" def fake_push(cfg, host, local_root, remote_root, log, *, models): pushed.append((str(local_root), remote_root)) log(f"push {Path(local_root).name}") return 2 monkeypatch.setattr("gpu_rent.provision.put_text", fake_put) monkeypatch.setattr("gpu_rent.provision.run_ssh", fake_ssh) monkeypatch.setattr("gpu_rent.sync_files.push_tree", fake_push) monkeypatch.setattr("gpu_rent.paths.assistent_extensions_dir", lambda: root) monkeypatch.setattr( "gpu_rent.paths.assistent_personas_dir", lambda: tmp_path / "no-personas" ) monkeypatch.setattr( "gpu_rent.paths.assistent_personas_manifest_path", lambda: tmp_path / "no.yaml" ) monkeypatch.setattr( "gpu_rent.provision.seed_civitai_examples", lambda *a, **k: None ) cfg = MagicMock() cfg.assistent_extensions_dir = str(root) logs: list[str] = [] seed_assistent_personas(cfg, "1.2.3.4", logs.append) assert len(pushed) == 1 assert pushed[0][1].endswith("/Assistent/extensions/sample-pack") assert "/mnt/swarm_data/Assistent/_base/assistant.json" in puts asst = json.loads(puts["/mnt/swarm_data/Assistent/_base/assistant.json"]) assert asst["default_persona"] == "neutral" assert asst["num_ctx"] == 16384 assert any("rm -f" in c and "personas.json" in c for c in ssh_cmds) assert any("packs/1" in m for m in logs) def test_seed_skips_empty_pack_dir(tmp_path: Path, monkeypatch): root = tmp_path / "assistent-extensions" (root / "empty").mkdir(parents=True) pushed: list[str] = [] def fake_push(cfg, host, local_root, remote_root, log, *, models): pushed.append(remote_root) return 0 monkeypatch.setattr("gpu_rent.provision.put_text", lambda *a, **k: None) monkeypatch.setattr( "gpu_rent.provision.run_ssh", lambda *a, **k: json.dumps({"vram_mib": 24576}) if ".gpu-rent-gpu" in str(a) else "", ) monkeypatch.setattr("gpu_rent.sync_files.push_tree", fake_push) monkeypatch.setattr("gpu_rent.paths.assistent_extensions_dir", lambda: root) monkeypatch.setattr( "gpu_rent.paths.assistent_personas_dir", lambda: tmp_path / "no-personas" ) monkeypatch.setattr( "gpu_rent.paths.assistent_personas_manifest_path", lambda: tmp_path / "no.yaml" ) monkeypatch.setattr( "gpu_rent.provision.seed_civitai_examples", lambda *a, **k: None ) cfg = MagicMock() cfg.assistent_extensions_dir = str(root) logs: list[str] = [] seed_assistent_personas(cfg, "1.2.3.4", logs.append) assert pushed == [] assert any("нет assistent-pack.yaml" in m for m in logs) assert any("packs/0" in m for m in logs)