351 lines
13 KiB
Python
351 lines
13 KiB
Python
import subprocess
|
|
import sys
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from click.exceptions import BadParameter
|
|
from typer.testing import CliRunner
|
|
|
|
from haiku.rag.cli import _cli as cli
|
|
from haiku.rag.cli import _parse_meta_options
|
|
from haiku.rag.cli import cli as cli_wrapper
|
|
from haiku.rag.store.exceptions import MigrationRequiredError
|
|
|
|
runner = CliRunner()
|
|
|
|
|
|
def test_importing_cli_does_not_load_heavy_dependencies():
|
|
"""Importing the CLI must not pull the heavy runtime dependencies; they cost
|
|
seconds of startup. Runs in a subprocess because another test in the same
|
|
session may already have imported them."""
|
|
result = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-c",
|
|
"import haiku.rag.cli, sys; "
|
|
"loaded = {'lancedb', 'pyarrow', 'pydantic_ai'} & sys.modules.keys(); "
|
|
"assert not loaded, loaded",
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
assert result.returncode == 0, result.stderr
|
|
|
|
|
|
class TestParseMetaOptions:
|
|
def test_empty_input(self):
|
|
assert _parse_meta_options(None) == {}
|
|
assert _parse_meta_options([]) == {}
|
|
|
|
def test_simple_key_value(self):
|
|
result = _parse_meta_options(["author=alice", "topic=notes"])
|
|
assert result == {"author": "alice", "topic": "notes"}
|
|
|
|
def test_missing_equals_raises(self):
|
|
with pytest.raises(BadParameter):
|
|
_parse_meta_options(["no_equals_here"])
|
|
|
|
def test_empty_key_raises(self):
|
|
with pytest.raises(BadParameter):
|
|
_parse_meta_options(["=value"])
|
|
|
|
def test_json_number(self):
|
|
result = _parse_meta_options(["version=3"])
|
|
assert result == {"version": 3}
|
|
assert isinstance(result["version"], int)
|
|
|
|
def test_json_float(self):
|
|
result = _parse_meta_options(["score=3.14"])
|
|
assert result == {"score": 3.14}
|
|
assert isinstance(result["score"], float)
|
|
|
|
def test_json_bool(self):
|
|
result = _parse_meta_options(["published=true", "draft=false"])
|
|
assert result == {"published": True, "draft": False}
|
|
|
|
def test_json_null(self):
|
|
result = _parse_meta_options(["empty=null"])
|
|
assert result == {"empty": None}
|
|
|
|
def test_json_array(self):
|
|
result = _parse_meta_options(['tags=["a","b","c"]'])
|
|
assert result == {"tags": ["a", "b", "c"]}
|
|
|
|
def test_json_object(self):
|
|
result = _parse_meta_options(['nested={"x": 1}'])
|
|
assert result == {"nested": {"x": 1}}
|
|
|
|
def test_plain_string_not_json(self):
|
|
result = _parse_meta_options(["name=hello world"])
|
|
assert result == {"name": "hello world"}
|
|
assert isinstance(result["name"], str)
|
|
|
|
def test_value_with_equals_sign(self):
|
|
result = _parse_meta_options(["equation=a=b+c"])
|
|
assert result == {"equation": "a=b+c"}
|
|
|
|
|
|
class TestRebuildValidation:
|
|
def test_embed_only_and_rechunk_mutually_exclusive(self):
|
|
result = runner.invoke(
|
|
cli, ["rebuild", "--embed-only", "--rechunk", "--db", "/tmp/fake.lancedb"]
|
|
)
|
|
assert result.exit_code == 1
|
|
assert "mutually exclusive" in result.output
|
|
|
|
|
|
class TestCliMigrationError:
|
|
def test_catches_migration_required_error(self):
|
|
with patch("haiku.rag.cli._cli") as mock_cli:
|
|
mock_cli.side_effect = MigrationRequiredError(
|
|
"Database requires migration. Run 'haiku-rag migrate' to upgrade."
|
|
)
|
|
|
|
with pytest.raises(SystemExit) as exc_info:
|
|
cli_wrapper()
|
|
assert exc_info.value.code == 1
|
|
|
|
|
|
class TestTagCommands:
|
|
def test_tag_round_trip(self, temp_db_path):
|
|
db = str(temp_db_path)
|
|
|
|
result = runner.invoke(cli, ["init", "--db", db])
|
|
assert result.exit_code == 0
|
|
|
|
result = runner.invoke(cli, ["tag", "create", "release-1", "--db", db])
|
|
assert result.exit_code == 0
|
|
assert "release-1" in result.output
|
|
|
|
result = runner.invoke(cli, ["tag", "list", "--db", db])
|
|
assert result.exit_code == 0
|
|
assert "release-1" in result.output
|
|
assert "partial" not in result.output
|
|
|
|
result = runner.invoke(cli, ["history", "--db", db, "-t", "documents"])
|
|
assert result.exit_code == 0
|
|
assert "release-1" in result.output
|
|
|
|
result = runner.invoke(cli, ["tag", "create", "release-1", "--db", db])
|
|
assert result.exit_code == 1
|
|
assert "already exists" in result.output
|
|
|
|
result = runner.invoke(cli, ["tag", "delete", "release-1", "--db", db])
|
|
assert result.exit_code == 0
|
|
|
|
result = runner.invoke(cli, ["tag", "list", "--db", db])
|
|
assert result.exit_code == 0
|
|
assert "No tags" in result.output
|
|
|
|
result = runner.invoke(cli, ["tag", "delete", "release-1", "--db", db])
|
|
assert result.exit_code == 1
|
|
assert "does not exist" in result.output
|
|
|
|
def test_tag_create_rejected_when_migrations_pending(self, temp_db_path):
|
|
"""A writable tag operation must hit the migration gate and must not
|
|
mutate a legacy database (e.g. by creating missing tables)."""
|
|
import asyncio
|
|
|
|
import lancedb
|
|
|
|
from haiku.rag.store.engine import Store
|
|
|
|
async def _prepare_legacy_db():
|
|
async with Store(temp_db_path, create=True) as store:
|
|
await store.set_haiku_version("0.19.0")
|
|
db = await lancedb.connect_async(temp_db_path.absolute())
|
|
await db.drop_table("document_meta")
|
|
db.close()
|
|
|
|
asyncio.run(_prepare_legacy_db())
|
|
|
|
result = runner.invoke(
|
|
cli, ["tag", "create", "release-1", "--db", str(temp_db_path)]
|
|
)
|
|
assert result.exit_code == 1
|
|
assert isinstance(result.exception, MigrationRequiredError)
|
|
|
|
async def _table_names() -> list[str]:
|
|
db = await lancedb.connect_async(temp_db_path.absolute())
|
|
tables = (await db.list_tables()).tables
|
|
db.close()
|
|
return tables
|
|
|
|
assert "document_meta" not in asyncio.run(_table_names())
|
|
|
|
def test_tag_commands_missing_database_exit_nonzero(self, tmp_path):
|
|
missing = str(tmp_path / "does_not_exist.lancedb")
|
|
for args in (
|
|
["tag", "create", "r1", "--db", missing],
|
|
["tag", "delete", "r1", "--db", missing],
|
|
["tag", "list", "--db", missing],
|
|
):
|
|
result = runner.invoke(cli, args)
|
|
assert result.exit_code == 1, args
|
|
assert "does not exist" in result.output, args
|
|
|
|
def test_tag_create_invalid_name_fails_cleanly(self, temp_db_path):
|
|
"""lance restricts ref names to alphanumeric, '.', '-', '_'; the CLI
|
|
surfaces that as a clean error instead of a traceback."""
|
|
db = str(temp_db_path)
|
|
result = runner.invoke(cli, ["init", "--db", db])
|
|
assert result.exit_code == 0
|
|
|
|
result = runner.invoke(cli, ["tag", "create", "[red]release[/red]", "--db", db])
|
|
assert result.exit_code == 1
|
|
assert "Error:" in result.output
|
|
assert "Ref characters" in result.output
|
|
|
|
|
|
class TestTagRestore:
|
|
def test_restore_requires_confirmation_and_decline_changes_nothing(
|
|
self, temp_db_path
|
|
):
|
|
db = str(temp_db_path)
|
|
assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0
|
|
assert runner.invoke(cli, ["tag", "create", "r1", "--db", db]).exit_code == 0
|
|
|
|
result = runner.invoke(cli, ["tag", "restore", "r1", "--db", db], input="n\n")
|
|
assert result.exit_code == 1
|
|
assert "live database state" in result.output
|
|
assert "Stop all ingestion" in result.output
|
|
assert "not transactionally atomic" in result.output
|
|
assert "safety tag" in result.output
|
|
|
|
result = runner.invoke(cli, ["tag", "list", "--db", db])
|
|
assert "before-restore" not in result.output
|
|
|
|
def test_restore_non_interactive_without_yes_fails(self, temp_db_path):
|
|
db = str(temp_db_path)
|
|
assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0
|
|
assert runner.invoke(cli, ["tag", "create", "r1", "--db", db]).exit_code == 0
|
|
|
|
result = runner.invoke(cli, ["tag", "restore", "r1", "--db", db])
|
|
assert result.exit_code == 1
|
|
|
|
result = runner.invoke(cli, ["tag", "list", "--db", db])
|
|
assert "before-restore" not in result.output
|
|
|
|
def test_restore_with_yes(self, temp_db_path):
|
|
db = str(temp_db_path)
|
|
assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0
|
|
assert runner.invoke(cli, ["tag", "create", "r1", "--db", db]).exit_code == 0
|
|
|
|
result = runner.invoke(cli, ["tag", "restore", "r1", "--yes", "--db", db])
|
|
assert result.exit_code == 0
|
|
assert "Restored database to tag 'r1'" in result.output
|
|
assert "before-restore-" in result.output
|
|
assert "now live" in result.output
|
|
assert "migrate" in result.output
|
|
|
|
result = runner.invoke(cli, ["tag", "list", "--db", db])
|
|
assert "before-restore-" in result.output
|
|
|
|
def test_restore_missing_tag_errors(self, temp_db_path):
|
|
db = str(temp_db_path)
|
|
assert runner.invoke(cli, ["init", "--db", db]).exit_code == 0
|
|
|
|
result = runner.invoke(cli, ["tag", "restore", "nope", "--yes", "--db", db])
|
|
assert result.exit_code == 1
|
|
assert "does not exist" in result.output
|
|
|
|
def test_restore_partial_tag_errors(self, temp_db_path):
|
|
import asyncio
|
|
|
|
from haiku.rag.store.engine import Store
|
|
|
|
async def _partial_tag():
|
|
async with Store(temp_db_path, create=True) as store:
|
|
version = await store.chunks_table.version()
|
|
await store.chunks_table.tags.create("stale", version)
|
|
|
|
asyncio.run(_partial_tag())
|
|
|
|
result = runner.invoke(
|
|
cli, ["tag", "restore", "stale", "--yes", "--db", str(temp_db_path)]
|
|
)
|
|
assert result.exit_code == 1
|
|
assert "partial" in result.output
|
|
assert "documents" in result.output
|
|
|
|
def test_restore_missing_database_exits_nonzero(self, tmp_path):
|
|
missing = tmp_path / "does_not_exist.lancedb"
|
|
result = runner.invoke(
|
|
cli, ["tag", "restore", "r1", "--yes", "--db", str(missing)]
|
|
)
|
|
assert result.exit_code == 1
|
|
assert "does not exist" in result.output
|
|
|
|
# Without --yes the missing database is reported before the
|
|
# confirmation prompt, not after the user confirms.
|
|
result = runner.invoke(cli, ["tag", "restore", "r1", "--db", str(missing)])
|
|
assert result.exit_code == 1
|
|
assert "does not exist" in result.output
|
|
assert "Continue?" not in result.output
|
|
|
|
def test_tag_help_includes_restore(self):
|
|
result = runner.invoke(cli, ["tag", "--help"])
|
|
assert result.exit_code == 0
|
|
assert "restore" in result.output
|
|
|
|
result = runner.invoke(cli, ["--help"])
|
|
assert "--before" not in result.output
|
|
assert "--at" not in result.output
|
|
|
|
|
|
class TestAskAnalyzeImageOption:
|
|
def test_ask_forwards_image_paths(self):
|
|
from unittest.mock import AsyncMock
|
|
|
|
from haiku.rag.app import HaikuRAGApp
|
|
|
|
with patch.object(HaikuRAGApp, "ask", new_callable=AsyncMock) as mock_ask:
|
|
result = runner.invoke(
|
|
cli,
|
|
["ask", "q", "--image", "/tmp/a.png", "--image", "/tmp/b.jpg"],
|
|
)
|
|
assert result.exit_code == 0
|
|
from pathlib import Path
|
|
|
|
assert mock_ask.call_args.kwargs["images"] == [
|
|
Path("/tmp/a.png"),
|
|
Path("/tmp/b.jpg"),
|
|
]
|
|
|
|
def test_analyze_forwards_image_paths(self):
|
|
from unittest.mock import AsyncMock
|
|
|
|
from haiku.rag.app import HaikuRAGApp
|
|
|
|
with patch.object(HaikuRAGApp, "analyze", new_callable=AsyncMock) as mock:
|
|
result = runner.invoke(cli, ["analyze", "q", "--image", "/tmp/a.png"])
|
|
assert result.exit_code == 0
|
|
from pathlib import Path
|
|
|
|
assert mock.call_args.kwargs["images"] == [Path("/tmp/a.png")]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_ask_reads_image_bytes(self, temp_db_path, tmp_path):
|
|
from io import BytesIO
|
|
from unittest.mock import AsyncMock
|
|
|
|
from PIL import Image as PILImage
|
|
|
|
from haiku.rag.app import HaikuRAGApp
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
buffer = BytesIO()
|
|
PILImage.new("RGB", (4, 4)).save(buffer, format="PNG")
|
|
img_path = tmp_path / "img.png"
|
|
img_path.write_bytes(buffer.getvalue())
|
|
|
|
async with HaikuRAG(temp_db_path, create=True):
|
|
pass
|
|
|
|
with patch.object(
|
|
HaikuRAG, "ask", new_callable=AsyncMock, return_value=("answer", [])
|
|
) as mock_ask:
|
|
app = HaikuRAGApp(db_path=temp_db_path)
|
|
await app.ask("q", images=[img_path])
|
|
|
|
assert mock_ask.call_args.kwargs["images"] == [buffer.getvalue()]
|