146 lines
4.8 KiB
Python
146 lines
4.8 KiB
Python
"""Exercise cli/llm.py's `index` and `serve` commands with mocked backends."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from typer.testing import CliRunner
|
|
|
|
from cli.llm import app
|
|
|
|
runner = CliRunner()
|
|
|
|
_STATS = {"indexed": 3, "skipped": 1, "chunks": 7}
|
|
|
|
|
|
class TestIndexComments:
|
|
@patch("llm.source.iter_comment_docs")
|
|
@patch("llm.pool.HostPool.from_config")
|
|
@patch("llm.index.index_docs")
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_default_collection_is_comments(
|
|
self, mock_load, mock_bib, mock_index_docs, mock_from_config, mock_iter
|
|
):
|
|
cfg = MagicMock()
|
|
mock_load.return_value = cfg
|
|
store = MagicMock()
|
|
mock_bib.return_value = store
|
|
docs = iter(["doc1", "doc2"])
|
|
mock_iter.return_value = docs
|
|
pool = MagicMock()
|
|
mock_from_config.return_value = pool
|
|
mock_index_docs.return_value = _STATS
|
|
|
|
result = runner.invoke(app, ["index"])
|
|
|
|
assert result.exit_code == 0
|
|
mock_iter.assert_called_once_with(store, docket="")
|
|
mock_index_docs.assert_called_once()
|
|
kwargs = mock_index_docs.call_args.kwargs
|
|
assert kwargs["collection"] == "comments"
|
|
assert kwargs["cfg"] is cfg
|
|
assert kwargs["pool"] is pool
|
|
assert kwargs["force"] is False
|
|
assert "indexed=3 skipped=1 chunks=7" in result.output
|
|
|
|
|
|
class TestIndexCorpus:
|
|
@patch("llm.source.iter_corpus_docs")
|
|
@patch("llm.pool.HostPool.from_config")
|
|
@patch("llm.index.index_docs")
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_corpus_collection(
|
|
self, mock_load, mock_bib, mock_index_docs, mock_from_config, mock_iter
|
|
):
|
|
cfg = MagicMock()
|
|
mock_load.return_value = cfg
|
|
store = MagicMock()
|
|
mock_bib.return_value = store
|
|
mock_iter.return_value = iter(["doc1"])
|
|
mock_from_config.return_value = MagicMock()
|
|
mock_index_docs.return_value = _STATS
|
|
|
|
result = runner.invoke(app, ["index", "--collection", "corpus"])
|
|
|
|
assert result.exit_code == 0
|
|
mock_iter.assert_called_once_with(store)
|
|
kwargs = mock_index_docs.call_args.kwargs
|
|
assert kwargs["collection"] == "corpus"
|
|
assert "indexed=3 skipped=1 chunks=7" in result.output
|
|
|
|
|
|
class TestIndexRules:
|
|
@patch("llm.source.iter_rule_docs")
|
|
@patch("llm.pool.HostPool.from_config")
|
|
@patch("llm.index.index_docs")
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_rules_collection(
|
|
self, mock_load, mock_bib, mock_index_docs, mock_from_config, mock_iter
|
|
):
|
|
cfg = MagicMock()
|
|
mock_load.return_value = cfg
|
|
store = MagicMock()
|
|
mock_bib.return_value = store
|
|
mock_iter.return_value = iter(["doc1"])
|
|
mock_from_config.return_value = MagicMock()
|
|
mock_index_docs.return_value = _STATS
|
|
|
|
result = runner.invoke(
|
|
app,
|
|
["index", "--collection", "rules", "--key", "abc123", "--key", "def456"],
|
|
)
|
|
|
|
assert result.exit_code == 0
|
|
mock_iter.assert_called_once_with(store, keys=("abc123", "def456"))
|
|
kwargs = mock_index_docs.call_args.kwargs
|
|
assert kwargs["collection"] == "rules"
|
|
assert "indexed=3 skipped=1 chunks=7" in result.output
|
|
|
|
|
|
class TestIndexBadCollection:
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_bad_collection_raises_bad_parameter(self, mock_load, mock_bib):
|
|
mock_load.return_value = MagicMock()
|
|
mock_bib.return_value = MagicMock()
|
|
|
|
result = runner.invoke(app, ["index", "--collection", "bogus"])
|
|
|
|
assert result.exit_code == 2
|
|
assert "collection must be 'comments', 'corpus' or 'rules'" in result.output
|
|
|
|
|
|
class TestIndexLimit:
|
|
@patch("llm.source.iter_comment_docs")
|
|
@patch("llm.pool.HostPool.from_config")
|
|
@patch("llm.index.index_docs")
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_limit_truncates_docs(
|
|
self, mock_load, mock_bib, mock_index_docs, mock_from_config, mock_iter
|
|
):
|
|
mock_load.return_value = MagicMock()
|
|
mock_bib.return_value = MagicMock()
|
|
mock_iter.return_value = iter([f"doc{i}" for i in range(5)])
|
|
mock_from_config.return_value = MagicMock()
|
|
mock_index_docs.return_value = _STATS
|
|
|
|
result = runner.invoke(app, ["index", "--limit", "2"])
|
|
|
|
assert result.exit_code == 0
|
|
docs_arg = mock_index_docs.call_args.args[0]
|
|
assert list(docs_arg) == ["doc0", "doc1"]
|
|
|
|
|
|
class TestServe:
|
|
@patch("uvicorn.run")
|
|
def test_serve_starts_uvicorn(self, mock_run):
|
|
result = runner.invoke(app, ["serve", "--host", "0.0.0.0", "--port", "9000"])
|
|
assert result.exit_code == 0
|
|
mock_run.assert_called_once_with(
|
|
"llm.api:app", host="0.0.0.0", port=9000, log_level="info"
|
|
)
|