Leonid Pershin
6cdd6ecfa1
- Enhanced the `with_token` function to support Gitea and Forgejo URLs by constructing the appropriate authentication format for generic HTTPS. - Updated tests to validate the new functionality, ensuring correct token injection for Gitea URLs while maintaining existing behavior for GitHub and GitLab. - Improved overall robustness of token handling across different platforms.
95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
from pathlib import Path
|
|
|
|
from gpu_rent.civitai import pick_primary_file
|
|
from gpu_rent.manifests import extract_version_id, repo_dirname, parse_models
|
|
from gpu_rent.payload import has_payload, model_push_set
|
|
from gpu_rent.remote.clone_ext import is_sha, strip_auth, with_token
|
|
|
|
|
|
def test_extract_version_id():
|
|
assert extract_version_id("https://civitai.red/models/1?modelVersionId=9") == 9
|
|
assert extract_version_id("https://civitai.com/api/download/models/77") == 77
|
|
|
|
|
|
def test_parse_models_fills_id_from_url(tmp_path: Path):
|
|
path = tmp_path / "m.yaml"
|
|
path.write_text("lora:\n - url: https://civitai.red/models/1?modelVersionId=42\n", encoding="utf-8")
|
|
entries = parse_models(path)
|
|
assert entries[0].version_id == 42
|
|
|
|
|
|
def test_repo_dirname():
|
|
from gpu_rent.manifests import GitRepo
|
|
|
|
repo = GitRepo(kind="swarmui", url="https://github.com/org/SwarmUI-Foo.git", ref="main", directory=None)
|
|
assert repo_dirname(repo) == "SwarmUI-Foo"
|
|
|
|
|
|
def test_pick_primary_file():
|
|
version = {
|
|
"files": [
|
|
{"name": "a.ckpt", "primary": False},
|
|
{"name": "b.safetensors", "primary": True},
|
|
]
|
|
}
|
|
assert pick_primary_file(version)["name"] == "b.safetensors"
|
|
|
|
|
|
def test_payload_pairs_sidecar_skips_orphan(tmp_path: Path):
|
|
root = tmp_path / "Models" / "Lora"
|
|
root.mkdir(parents=True)
|
|
(root / ".gitkeep").write_text("x", encoding="utf-8")
|
|
(root / "README.md").write_text("x", encoding="utf-8")
|
|
(root / "foo.safetensors").write_bytes(b"weights")
|
|
(root / "foo.civitai.json").write_text("{}", encoding="utf-8")
|
|
(root / "orphan.json").write_text("{}", encoding="utf-8")
|
|
names = {p.name for p in model_push_set(tmp_path / "Models")}
|
|
assert names == {"foo.safetensors", "foo.civitai.json"}
|
|
empty = tmp_path / "Wildcards"
|
|
empty.mkdir()
|
|
(empty / ".gitkeep").write_text("", encoding="utf-8")
|
|
assert not has_payload(empty)
|
|
|
|
|
|
def test_git_token_injection():
|
|
url = "https://github.com/org/Ext.git"
|
|
assert "x-access-token:secret@" in with_token(url, "secret")
|
|
assert strip_auth(with_token(url, "secret")) == url
|
|
gitea = "https://gitea.hsrv.site/mrleo1nid/swarm-assistent.git"
|
|
authed = with_token(gitea, "secret")
|
|
assert "oauth2:secret@gitea.hsrv.site" in authed
|
|
assert strip_auth(authed) == gitea
|
|
assert is_sha("a" * 40)
|
|
assert not is_sha("main")
|
|
|
|
|
|
def test_scrub_origin(tmp_path: Path, monkeypatch):
|
|
from gpu_rent.remote import clone_ext
|
|
|
|
dest = tmp_path / "repo"
|
|
dest.mkdir()
|
|
(dest / ".git").mkdir()
|
|
calls: list[list[str]] = []
|
|
|
|
def fake_out(argv, cwd=None):
|
|
if "get-url" in argv:
|
|
return "https://x-access-token:secret@github.com/org/Ext.git"
|
|
raise AssertionError(argv)
|
|
|
|
def fake_run(argv, cwd=None):
|
|
calls.append(argv)
|
|
|
|
monkeypatch.setattr(clone_ext, "out", fake_out)
|
|
monkeypatch.setattr(clone_ext, "run", fake_run)
|
|
clone_ext.scrub_origin(dest, "https://github.com/org/Ext.git")
|
|
assert calls == [
|
|
[
|
|
"git",
|
|
"-C",
|
|
str(dest),
|
|
"remote",
|
|
"set-url",
|
|
"origin",
|
|
"https://github.com/org/Ext.git",
|
|
]
|
|
] |