Leonid PershinandCursor
07bb521b70
Introduce books/ with civitai-krea2 and HF fictext builders, sha-diff seed to VM, and drop train.jsonl from the civitai scrape pipeline. Co-authored-by: Cursor <cursoragent@cursor.com>
130 lines
4.2 KiB
Python
130 lines
4.2 KiB
Python
"""Tests for Assistent books build + seed helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
|
|
from gpu_rent.books import (
|
|
BOOK_CIVITAI,
|
|
build_civitai_book,
|
|
civitai_search_to_book_row,
|
|
fictext_row_to_book_rows,
|
|
read_meta,
|
|
sha256_file,
|
|
write_meta,
|
|
)
|
|
from gpu_rent.provision import seed_books
|
|
|
|
|
|
def test_civitai_search_to_book_row():
|
|
row = {
|
|
"id": 1,
|
|
"prompt": "woman, redhead",
|
|
"negative": "blur",
|
|
"tags": ["woman"],
|
|
"rating": "pg13",
|
|
"params": {"steps": 8, "cfg": 1},
|
|
"loras": [{"versionId": 9, "weight": 0.7}],
|
|
"modelVersionId": 100,
|
|
"kind": "checkpoint",
|
|
"score": 10,
|
|
}
|
|
out = civitai_search_to_book_row(row)
|
|
assert out["book"] == BOOK_CIVITAI
|
|
assert "woman" in out["text"]
|
|
assert out["meta"]["loras"][0]["versionId"] == 9
|
|
|
|
|
|
def test_fictext_chunks_skip_minor():
|
|
rows = fictext_row_to_book_rows(
|
|
{"title": "T", "tags": ["loli"], "text": "loli content here"},
|
|
"ru-fictext-rplus",
|
|
row_index=0,
|
|
)
|
|
assert rows == []
|
|
|
|
|
|
def test_fictext_chunks_ok():
|
|
text = "First paragraph.\n\nSecond paragraph with more words."
|
|
rows = fictext_row_to_book_rows(
|
|
{"title": "Story", "tags": ["romance"], "text": text},
|
|
"ru-fictext-rplus",
|
|
row_index=3,
|
|
)
|
|
assert len(rows) >= 1
|
|
assert rows[0]["title"] == "Story"
|
|
assert "romance" in rows[0]["tags"]
|
|
|
|
|
|
def test_build_civitai_book(tmp_path: Path):
|
|
ds = tmp_path / "datasets" / "civitai"
|
|
ds.mkdir(parents=True)
|
|
search_row = {
|
|
"id": 5,
|
|
"prompt": "test prompt",
|
|
"negative": "",
|
|
"tags": ["a"],
|
|
"rating": "pg",
|
|
"params": {},
|
|
"loras": [],
|
|
}
|
|
(ds / "search.jsonl").write_text(json.dumps(search_row) + "\n", encoding="utf-8")
|
|
counts = build_civitai_book(out_root=tmp_path, log=lambda m: None)
|
|
assert counts["rows"] == 1
|
|
meta = read_meta(BOOK_CIVITAI, tmp_path)
|
|
assert meta and meta.get("row_count") == 1
|
|
book_search = tmp_path / "books" / BOOK_CIVITAI / "search.jsonl"
|
|
assert book_search.is_file()
|
|
assert sha256_file(book_search) == meta["content_sha"]
|
|
|
|
|
|
def test_seed_books_skips_unchanged(tmp_path: Path, monkeypatch):
|
|
book = tmp_path / "books" / BOOK_CIVITAI
|
|
book.mkdir(parents=True)
|
|
(book / "search.jsonl").write_text('{"id":"x"}\n', encoding="utf-8")
|
|
write_meta(BOOK_CIVITAI, row_count=1, root=tmp_path)
|
|
meta = read_meta(BOOK_CIVITAI, tmp_path)
|
|
sidecar = json.dumps({"content_sha": meta["content_sha"]})
|
|
|
|
monkeypatch.setattr("gpu_rent.paths.app_root", lambda: tmp_path)
|
|
monkeypatch.setattr(
|
|
"gpu_rent.provision.run_ssh",
|
|
lambda *a, **k: sidecar if ".gpu-rent-meta" in str(a[2]) else "",
|
|
)
|
|
put_calls: list[str] = []
|
|
monkeypatch.setattr(
|
|
"gpu_rent.ssh_ops.put_file",
|
|
lambda *a, **k: put_calls.append(str(a[3])),
|
|
)
|
|
monkeypatch.setattr("gpu_rent.ssh_ops.put_text", lambda *a, **k: None)
|
|
monkeypatch.setattr("gpu_rent.provision.seed_civitai_examples", lambda *a, **k: None)
|
|
|
|
logs: list[str] = []
|
|
seed_books(MagicMock(app_root=str(tmp_path)), "host", logs.append)
|
|
assert put_calls == []
|
|
assert any("unchanged" in m for m in logs)
|
|
|
|
|
|
def test_seed_books_pushes_when_sha_diff(tmp_path: Path, monkeypatch):
|
|
book = tmp_path / "books" / BOOK_CIVITAI
|
|
book.mkdir(parents=True)
|
|
(book / "book.yaml").write_text("kind: book\n", encoding="utf-8")
|
|
(book / "search.jsonl").write_text('{"id":"x"}\n', encoding="utf-8")
|
|
write_meta(BOOK_CIVITAI, row_count=1, root=tmp_path)
|
|
|
|
monkeypatch.setattr("gpu_rent.paths.app_root", lambda: tmp_path)
|
|
monkeypatch.setattr("gpu_rent.provision.run_ssh", lambda *a, **k: "")
|
|
put_calls: list[str] = []
|
|
monkeypatch.setattr(
|
|
"gpu_rent.ssh_ops.put_file",
|
|
lambda *a, **k: put_calls.append(str(a[3])),
|
|
)
|
|
monkeypatch.setattr("gpu_rent.ssh_ops.put_text", lambda *a, **k: None)
|
|
|
|
logs: list[str] = []
|
|
seed_books(MagicMock(app_root=str(tmp_path)), "host", logs.append)
|
|
assert any("search.jsonl" in p for p in put_calls)
|
|
assert any("books/civitai-krea2" in m for m in logs)
|