49a2f6f723
prisma.fleet.run_pool owns the multi-host thread-pool loop that screen had inline: at most --per-host in-flight requests per Ollama host, LLM calls only on the pool, every SQLite touch and per-item commit on the calling thread, and a host that fails twice in a row is retired for the rest of the run (its queued items go to the surviving hosts). screen, eligibility and extract each supply prepare/apply/on_error glue; items with no retrievable full text resolve without a fleet call. CLI: prisma eligible/extract gain --hosts/--fleet-model/--num-ctx/--per-host matching prisma screen. Docs regenerated.
405 lines
13 KiB
Python
405 lines
13 KiB
Python
"""Deep exercising tests for cli/prisma.py — covers function bodies."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from typer.testing import CliRunner
|
|
|
|
from cli.prisma import app
|
|
|
|
runner = CliRunner()
|
|
|
|
|
|
def _mock_db():
|
|
db = MagicMock()
|
|
db.__enter__ = MagicMock(return_value=db)
|
|
db.__exit__ = MagicMock(return_value=False)
|
|
return db
|
|
|
|
|
|
class TestHold:
|
|
@patch("cli.prisma.subprocess.run")
|
|
def test_hold_true(self, mc_sub):
|
|
from cli.prisma import _hold
|
|
|
|
result = _hold(lambda: 42, hold=True)
|
|
assert result == 42
|
|
assert mc_sub.call_count == 2
|
|
|
|
def test_hold_false(self):
|
|
from cli.prisma import _hold
|
|
|
|
result = _hold(lambda: 99, hold=False)
|
|
assert result == 99
|
|
|
|
|
|
class TestInit:
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.project.init")
|
|
def test_basic(self, mc_init, mc_db_cls, mc_sub):
|
|
project = MagicMock()
|
|
project.name = "test-proj"
|
|
project.criteria = "criteria text"
|
|
project.extraction_template = "template"
|
|
project.reasons = "reasons"
|
|
mc_init.return_value = project
|
|
mc_db_cls.return_value = _mock_db()
|
|
|
|
result = runner.invoke(app, ["init", "test-proj", "--no-hold"])
|
|
assert result.exit_code == 0
|
|
assert "test-proj" in result.output
|
|
|
|
|
|
class TestExport:
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.export.load_item")
|
|
@patch("prisma.export.to_markdown", return_value="# Exported markdown")
|
|
def test_basic(self, mc_md, mc_load, mc_db_cls, mc_sub):
|
|
mc_db_cls.return_value = _mock_db()
|
|
mc_load.return_value = MagicMock()
|
|
|
|
result = runner.invoke(app, ["export", "123", "--no-hold"])
|
|
assert result.exit_code == 0
|
|
assert "Exported markdown" in result.output
|
|
|
|
|
|
class TestScreen:
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.llm.make_provider")
|
|
@patch("prisma.project.load")
|
|
@patch("prisma.screen.run", return_value={"screened": 10, "included": 7})
|
|
def test_basic(self, mc_run, mc_load, mc_prov, mc_db_cls, mc_sub):
|
|
mc_db_cls.return_value = _mock_db()
|
|
mc_prov.return_value = MagicMock()
|
|
mc_load.return_value = MagicMock()
|
|
|
|
result = runner.invoke(app, ["screen", "test-proj", "--no-hold"])
|
|
assert result.exit_code == 0
|
|
|
|
|
|
class TestEligible:
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.llm.make_provider")
|
|
@patch("prisma.project.load")
|
|
@patch("prisma.eligibility.run", return_value={"eligible": 5})
|
|
def test_basic(self, mc_run, mc_load, mc_prov, mc_db_cls, mc_sub):
|
|
mc_db_cls.return_value = _mock_db()
|
|
mc_prov.return_value = MagicMock()
|
|
mc_load.return_value = MagicMock()
|
|
|
|
result = runner.invoke(app, ["eligible", "test-proj", "--no-hold"])
|
|
assert result.exit_code == 0
|
|
mc_run.assert_called_once()
|
|
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.llm.make_provider")
|
|
@patch("prisma.project.load")
|
|
@patch("prisma.eligibility.run_fleet", return_value={"assessed": 5})
|
|
@patch("prisma.eligibility.run", return_value={"eligible": 5})
|
|
def test_fleet_flags_wired_to_driver(
|
|
self, mc_run, mc_fleet, mc_load, mc_prov, mc_db_cls, mc_sub
|
|
):
|
|
mc_db_cls.return_value = _mock_db()
|
|
mc_prov.return_value = MagicMock()
|
|
project = MagicMock()
|
|
mc_load.return_value = project
|
|
|
|
result = runner.invoke(
|
|
app,
|
|
[
|
|
"eligible",
|
|
"test-proj",
|
|
"--no-hold",
|
|
"--hosts",
|
|
"http://h1:11434, http://h2:11434",
|
|
"--fleet-model",
|
|
"qwen2.5:32b",
|
|
"--num-ctx",
|
|
"16384",
|
|
"--per-host",
|
|
"3",
|
|
"--max-chars",
|
|
"40000",
|
|
],
|
|
)
|
|
assert result.exit_code == 0
|
|
mc_run.assert_not_called()
|
|
mc_fleet.assert_called_once()
|
|
_, kwargs = mc_fleet.call_args
|
|
assert kwargs["hosts"] == ["http://h1:11434", "http://h2:11434"]
|
|
assert kwargs["model"] == "qwen2.5:32b"
|
|
assert kwargs["num_ctx"] == 16384
|
|
assert kwargs["per_host"] == 3
|
|
assert kwargs["max_fulltext_chars"] == 40000
|
|
|
|
|
|
class TestExtract:
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.llm.make_provider")
|
|
@patch("prisma.project.load")
|
|
@patch("prisma.extract.run", return_value={"extracted": 3})
|
|
def test_basic(self, mc_run, mc_load, mc_prov, mc_db_cls, mc_sub):
|
|
mc_db_cls.return_value = _mock_db()
|
|
mc_prov.return_value = MagicMock()
|
|
mc_load.return_value = MagicMock()
|
|
|
|
result = runner.invoke(app, ["extract", "test-proj", "--no-hold"])
|
|
assert result.exit_code == 0
|
|
mc_run.assert_called_once()
|
|
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.llm.make_provider")
|
|
@patch("prisma.project.load")
|
|
@patch("prisma.extract.run_fleet", return_value={"extracted": 3})
|
|
@patch("prisma.extract.run", return_value={"extracted": 3})
|
|
def test_fleet_flags_wired_to_driver(
|
|
self, mc_run, mc_fleet, mc_load, mc_prov, mc_db_cls, mc_sub
|
|
):
|
|
mc_db_cls.return_value = _mock_db()
|
|
mc_prov.return_value = MagicMock()
|
|
mc_load.return_value = MagicMock()
|
|
|
|
result = runner.invoke(
|
|
app,
|
|
[
|
|
"extract",
|
|
"test-proj",
|
|
"--no-hold",
|
|
"--hosts",
|
|
"http://h1:11434,http://h2:11434",
|
|
"--fleet-model",
|
|
"qwen2.5:32b",
|
|
"--num-ctx",
|
|
"16384",
|
|
"--per-host",
|
|
"2",
|
|
],
|
|
)
|
|
assert result.exit_code == 0
|
|
mc_run.assert_not_called()
|
|
mc_fleet.assert_called_once()
|
|
_, kwargs = mc_fleet.call_args
|
|
assert kwargs["hosts"] == ["http://h1:11434", "http://h2:11434"]
|
|
assert kwargs["model"] == "qwen2.5:32b"
|
|
assert kwargs["num_ctx"] == 16384
|
|
assert kwargs["per_host"] == 2
|
|
|
|
|
|
class TestFlow:
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.flow.count")
|
|
@patch("prisma.flow.mermaid", return_value="graph TD; A-->B")
|
|
def test_mermaid(self, mc_mermaid, mc_count, mc_db_cls, mc_sub):
|
|
mc_db_cls.return_value = _mock_db()
|
|
mc_count.return_value = MagicMock()
|
|
|
|
result = runner.invoke(app, ["flow", "test-proj", "--no-hold"])
|
|
assert result.exit_code == 0
|
|
assert "graph" in result.output
|
|
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.flow.count")
|
|
@patch("prisma.flow.text_summary", return_value="Identified: 100")
|
|
def test_text(self, mc_text, mc_count, mc_db_cls, mc_sub):
|
|
mc_db_cls.return_value = _mock_db()
|
|
mc_count.return_value = MagicMock()
|
|
|
|
result = runner.invoke(app, ["flow", "test-proj", "--text", "--no-hold"])
|
|
assert result.exit_code == 0
|
|
assert "Identified" in result.output
|
|
|
|
|
|
class TestFetch:
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.fetch.run", return_value={"fetched": 10, "failed": 2})
|
|
@patch("prisma.vpn.status", return_value={"status": "down"})
|
|
def test_no_vpn(self, mc_vpn_status, mc_fetch, mc_db_cls, mc_sub):
|
|
mc_db_cls.return_value = _mock_db()
|
|
|
|
result = runner.invoke(app, ["fetch", "test-proj", "--no-hold", "--no-vpn"])
|
|
assert result.exit_code == 0
|
|
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.fetch.run", return_value={"fetched": 10})
|
|
@patch("prisma.vpn.status", return_value={"status": "up"})
|
|
@patch("prisma.vpn.active")
|
|
@patch("prisma.vpn.verify_egress", return_value={"ip": "1.2.3.4", "country": "NL"})
|
|
def test_with_vpn(
|
|
self, mc_egress, mc_active, mc_status, mc_fetch, mc_db_cls, mc_sub
|
|
):
|
|
mc_db_cls.return_value = _mock_db()
|
|
ctx = MagicMock()
|
|
ctx.__enter__ = MagicMock(return_value="socks5://localhost:1080")
|
|
ctx.__exit__ = MagicMock(return_value=False)
|
|
mc_active.return_value = ctx
|
|
|
|
result = runner.invoke(app, ["fetch", "test-proj", "--no-hold"])
|
|
assert result.exit_code == 0
|
|
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.fetch.run", return_value={"fetched": 5})
|
|
@patch("prisma.vpn.status", return_value={"status": "down"})
|
|
def test_explicit_proxy(self, mc_status, mc_fetch, mc_db_cls, mc_sub):
|
|
mc_db_cls.return_value = _mock_db()
|
|
|
|
result = runner.invoke(
|
|
app,
|
|
["fetch", "test-proj", "--no-hold", "--proxy", "socks5://host:1080"],
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
|
|
class TestRunAll:
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.llm.make_provider")
|
|
@patch("prisma.project.load")
|
|
@patch("prisma.screen.run", return_value={"screened": 10})
|
|
@patch("prisma.fetch.run", return_value={"fetched": 5})
|
|
@patch("prisma.eligibility.run", return_value={"eligible": 3})
|
|
@patch("prisma.extract.run", return_value={"extracted": 2})
|
|
@patch("prisma.flow.count")
|
|
@patch("prisma.flow.text_summary", return_value="Summary: OK")
|
|
@patch("prisma.vpn.status", return_value={"status": "down"})
|
|
def test_full(
|
|
self,
|
|
mc_vpn,
|
|
mc_summary,
|
|
mc_count,
|
|
mc_extract,
|
|
mc_elig,
|
|
mc_fetch,
|
|
mc_screen,
|
|
mc_load,
|
|
mc_prov,
|
|
mc_db_cls,
|
|
mc_sub,
|
|
):
|
|
mc_db_cls.return_value = _mock_db()
|
|
mc_prov.return_value = MagicMock()
|
|
mc_load.return_value = MagicMock()
|
|
mc_count.return_value = MagicMock()
|
|
|
|
result = runner.invoke(app, ["run", "test-proj", "--no-hold"])
|
|
assert result.exit_code == 0
|
|
|
|
@patch("cli.prisma.subprocess.run")
|
|
@patch("zot.db.Db")
|
|
@patch("prisma.llm.make_provider")
|
|
@patch("prisma.project.load")
|
|
@patch("prisma.screen.run", return_value={"screened": 10})
|
|
@patch("prisma.eligibility.run", return_value={"eligible": 3})
|
|
@patch("prisma.flow.count")
|
|
@patch("prisma.flow.text_summary", return_value="Summary: OK")
|
|
@patch("prisma.vpn.status", return_value={"status": "down"})
|
|
def test_skip_fetch_extract(
|
|
self,
|
|
mc_vpn,
|
|
mc_summary,
|
|
mc_count,
|
|
mc_elig,
|
|
mc_screen,
|
|
mc_load,
|
|
mc_prov,
|
|
mc_db_cls,
|
|
mc_sub,
|
|
):
|
|
mc_db_cls.return_value = _mock_db()
|
|
mc_prov.return_value = MagicMock()
|
|
mc_load.return_value = MagicMock()
|
|
mc_count.return_value = MagicMock()
|
|
|
|
result = runner.invoke(
|
|
app, ["run", "test-proj", "--no-hold", "--skip-fetch", "--skip-extract"]
|
|
)
|
|
assert result.exit_code == 0
|
|
|
|
|
|
class TestVpnUp:
|
|
@patch("prisma.vpn.up")
|
|
def test_basic(self, mc_up):
|
|
mc_up.return_value = {
|
|
"name": "stack-prisma-vpn",
|
|
"droplet_id": 123,
|
|
"region": "nyc3",
|
|
"public_ip": "1.2.3.4",
|
|
"ssh_key": "~/.ssh/id_ed25519",
|
|
"proxy_url": "socks5://localhost:1080",
|
|
"sidecar": "fetch-proxy",
|
|
"zotero_proxy": "socks5://fetch-proxy:1080",
|
|
}
|
|
|
|
result = runner.invoke(app, ["vpn", "up"])
|
|
assert result.exit_code == 0
|
|
assert "1.2.3.4" in result.output
|
|
|
|
|
|
class TestVpnDown:
|
|
@patch("prisma.vpn.down", return_value={"destroyed": True})
|
|
def test_basic(self, mc_down):
|
|
result = runner.invoke(app, ["vpn", "down"])
|
|
assert result.exit_code == 0
|
|
|
|
|
|
class TestVpnStatus:
|
|
@patch("prisma.vpn.status", return_value={"status": "up", "ip": "1.2.3.4"})
|
|
def test_basic(self, mc_status):
|
|
result = runner.invoke(app, ["vpn", "status"])
|
|
assert result.exit_code == 0
|
|
|
|
|
|
class TestVpnAttachZotero:
|
|
@patch("prisma.vpn.attach_zotero_proxy")
|
|
@patch("prisma.vpn._DROPLET_JSON")
|
|
def test_basic(self, mc_json, mc_attach):
|
|
mc_json.is_file.return_value = True
|
|
mc_json.read_text.return_value = '{"public_ip": "1.2.3.4"}'
|
|
result = runner.invoke(app, ["vpn", "attach-zotero"])
|
|
assert result.exit_code == 0
|
|
|
|
@patch("prisma.vpn._DROPLET_JSON")
|
|
def test_no_droplet(self, mc_json):
|
|
mc_json.is_file.return_value = False
|
|
result = runner.invoke(app, ["vpn", "attach-zotero"])
|
|
assert result.exit_code != 0
|
|
|
|
|
|
class TestVpnDetachZotero:
|
|
@patch("prisma.vpn.detach_zotero_proxy")
|
|
def test_basic(self, mc_detach):
|
|
result = runner.invoke(app, ["vpn", "detach-zotero"])
|
|
assert result.exit_code == 0
|
|
|
|
|
|
class TestVpnVerify:
|
|
@patch("prisma.vpn.active")
|
|
@patch("prisma.vpn.verify_egress")
|
|
def test_basic(self, mc_egress, mc_active):
|
|
ctx = MagicMock()
|
|
ctx.__enter__ = MagicMock(return_value="socks5://localhost:1080")
|
|
ctx.__exit__ = MagicMock(return_value=False)
|
|
mc_active.return_value = ctx
|
|
mc_egress.return_value = {
|
|
"ip": "1.2.3.4",
|
|
"country": "Netherlands",
|
|
"country_iso": "NL",
|
|
"asn_org": "DigitalOcean",
|
|
}
|
|
|
|
result = runner.invoke(app, ["vpn", "verify"])
|
|
assert result.exit_code == 0
|
|
assert "Netherlands" in result.output
|