diff --git a/src/cli/main.py b/src/cli/main.py index 2cc2f91..7c97c61 100644 --- a/src/cli/main.py +++ b/src/cli/main.py @@ -222,5 +222,30 @@ def stats( typer.echo(f" - {s}") +@app.command(help="导出 collection 数据为 JSON 或 CSV.") +def export( + output: Annotated[str, typer.Option("--output", "-o", help="输出文件路径")], + fmt: Annotated[ + str, typer.Option("--format", "-f", help="导出格式: json | csv") + ] = "json", + config: ConfigOpt = DEFAULT_CONFIG_PATH, + collection: CollectionOpt = None, +): + """导出 collection 数据.""" + _init_config(config) + state = get_state() + searcher = state.get_searcher(_resolve_collection(collection)) + + if fmt == "json": + searcher.export_json(file_path=output) + elif fmt == "csv": + searcher.export_csv(file_path=output) + else: + typer.echo(f"错误: 不支持的格式 '{fmt}',可选: json, csv", err=True) + raise typer.Exit(code=1) + + typer.echo(f"[OK] 已导出到: {output}") + + if __name__ == "__main__": app() diff --git a/tests/test_export.py b/tests/test_export.py new file mode 100644 index 0000000..fc053cd --- /dev/null +++ b/tests/test_export.py @@ -0,0 +1,106 @@ +"""导出功能测试.""" +import csv +import io +import json + +import pytest + +from src.core.search import Searcher + + +class FakeExportCollection: + def count(self): + return 2 + + def get(self, include=None): + return { + "ids": ["doc_0", "doc_1"], + "documents": ["内容A。\n\n段落B。", "内容C。"], + "metadatas": [ + {"source_file": "a.md", "section_title": "标题A", "heading_level": 1, "chunk_index": 0}, + {"source_file": "b.md", "section_title": "", "heading_level": 0, "chunk_index": 0}, + ], + } + + +class FakeExportDB: + def get_or_create_collection(self, name): + return FakeExportCollection() + + def list_collections(self): + return [] + + +class FakeExportEmbedder: + @property + def dimension(self): + return 4 + + def embed(self, texts): + return [[0.1, 0.2, 0.3, 0.4] for _ in texts] + + +@pytest.fixture +def searcher(): + db = FakeExportDB() + embedder = FakeExportEmbedder() + return Searcher(db, embedder, "test_export") + + +def test_export_json_returns_valid_json(searcher): + """export_json 返回合法的 JSON 字符串.""" + output = searcher.export_json() + data = json.loads(output) + assert isinstance(data, list) + assert len(data) == 2 + + +def test_export_json_contains_all_fields(searcher): + """导出包含所有必要字段.""" + output = searcher.export_json() + data = json.loads(output) + first = data[0] + assert "id" in first + assert "content" in first + assert "source_file" in first + assert "section_title" in first + + +def test_export_csv_returns_valid_csv(searcher): + """export_csv 返回合法的 CSV 字符串.""" + output = searcher.export_csv() + reader = csv.DictReader(io.StringIO(output)) + rows = list(reader) + assert len(rows) == 2 + + +def test_export_csv_has_header(searcher): + """CSV 包含表头.""" + output = searcher.export_csv() + reader = csv.DictReader(io.StringIO(output)) + assert reader.fieldnames is not None + assert "content" in reader.fieldnames + assert "source_file" in reader.fieldnames + + +def test_export_empty_collection(): + """空 collection 导出空列表.""" + + class EmptyCollection: + def count(self): + return 0 + def get(self, include=None): + return {"ids": [], "documents": [], "metadatas": []} + + class EmptyDB: + def get_or_create_collection(self, name): + return EmptyCollection() + def list_collections(self): + return [] + + s = Searcher(EmptyDB(), FakeExportEmbedder(), "empty") + json_out = s.export_json() + assert json.loads(json_out) == [] + csv_out = s.export_csv() + lines = csv_out.strip().split("\n") + assert len(lines) == 1 # 仅表头