Enhance CLI and documentation for capturing VM inventory

- Introduced new `capture` commands in the CLI to allow users to merge VM inventory into local manifests without downloading weights. - Updated `README.md` and `cli.md` to include detailed instructions for the new capture functionality, including options for models and extensions. - Enhanced `decisions.md` to clarify the role of captured links in the manifest files. - Improved `extensions.md` to document the process of capturing installed extensions back to the local configuration. - Added new functions in `civitai.py` to support fetching model versions by hash and generating canonical URLs for models.
2026-08-21 05:56:52 +03:00
parent 71f4e4c2e3
commit 603165a4ba
10 changed files with 940 additions and 3 deletions
@@ -0,0 +1,111 @@
+"""Unit tests for capture merge / URL builders (no SSH)."""
+
+from pathlib import Path
+
+from gpu_rent.capture import (
+    ExtCaptureItem,
+    ModelCaptureItem,
+    merge_extensions_yaml,
+    merge_models_yaml,
+    resolve_model_item,
+    strip_git_auth,
+)
+from gpu_rent.civitai import civitai_model_url, version_ids_from_payload
+from gpu_rent.manifests import parse_extensions, parse_models
+
+
+def test_civitai_model_url():
+    assert (
+        civitai_model_url(2786499, 3231611, "civitai.red")
+        == "https://civitai.red/models/2786499?modelVersionId=3231611"
+    )
+
+
+def test_version_ids_from_payload():
+    vid, mid = version_ids_from_payload({"id": 9, "modelId": 1})
+    assert vid == 9 and mid == 1
+    vid2, mid2 = version_ids_from_payload({"id": 3, "model": {"id": 7}})
+    assert vid2 == 3 and mid2 == 7
+
+
+def test_resolve_from_sidecar():
+    item = resolve_model_item(
+        {
+            "kind": "lora",
+            "rel": "Lora/foo.safetensors",
+            "name": "foo.safetensors",
+            "version_id": 3107521,
+            "model_id": 2187487,
+            "sha256": "a" * 64,
+        },
+        token="",
+        api_host="civitai.red",
+        link_host="civitai.red",
+    )
+    assert item is not None
+    assert item.version_id == 3107521
+    assert "modelVersionId=3107521" in item.url
+
+
+def test_merge_models_dedupe(tmp_path: Path):
+    path = tmp_path / "models.yaml"
+    path.write_text(
+        "lora:\n  - url: https://civitai.red/models/1?modelVersionId=100\n",
+        encoding="utf-8",
+    )
+    items = [
+        ModelCaptureItem("lora", 100, 1, "https://civitai.red/models/1?modelVersionId=100", "old"),
+        ModelCaptureItem("lora", 200, 2, "https://civitai.red/models/2?modelVersionId=200", "new"),
+        ModelCaptureItem("checkpoint", 300, 3, "https://civitai.red/models/3?modelVersionId=300", "ckpt"),
+    ]
+    added, skipped = merge_models_yaml(path, items, dry_run=False)
+    assert len(added) == 2
+    assert any("100" in s for s in skipped)
+    entries = parse_models(path)
+    vids = {e.version_id for e in entries}
+    assert vids == {100, 200, 300}
+    assert (tmp_path / "models.yaml.bak").is_file()
+
+
+def test_merge_models_dry_run(tmp_path: Path):
+    path = tmp_path / "models.yaml"
+    path.write_text("lora: []\n", encoding="utf-8")
+    before = path.read_text(encoding="utf-8")
+    items = [
+        ModelCaptureItem("lora", 1, 1, "https://civitai.red/models/1?modelVersionId=1", "x"),
+    ]
+    added, _ = merge_models_yaml(path, items, dry_run=True)
+    assert len(added) == 1
+    assert path.read_text(encoding="utf-8") == before
+
+
+def test_merge_extensions_dedupe(tmp_path: Path):
+    path = tmp_path / "extensions.yaml"
+    path.write_text(
+        "swarmui:\n  - url: https://github.com/org/A.git\n    ref: main\n    dir: A\n",
+        encoding="utf-8",
+    )
+    items = [
+        ExtCaptureItem("swarmui", "https://github.com/org/A.git", "main", "A"),
+        ExtCaptureItem("comfy", "https://github.com/org/B.git", "v1", "B"),
+        ExtCaptureItem(
+            "swarmui",
+            "https://x-access-token:SECRET@github.com/org/C.git",
+            "main",
+            "C",
+        ),
+    ]
+    added, skipped = merge_extensions_yaml(path, items, dry_run=False)
+    assert len(added) == 2
+    assert len(skipped) == 1
+    repos = parse_extensions(path)
+    urls = [r.url for r in repos]
+    assert "https://github.com/org/B.git" in urls
+    assert all("SECRET" not in u for u in urls)
+
+
+def test_strip_git_auth():
+    assert (
+        strip_git_auth("https://x-access-token:tok@github.com/org/r.git")
+        == "https://github.com/org/r.git"
+    )