from __future__ import annotations import io import unittest from collections.abc import Callable from contextlib import redirect_stdout from pathlib import Path from typing import TypeVar from unittest.mock import MagicMock, patch from chromy.handlers.count_collection import handle_count_collection from chromy.handlers.create_collection import handle_create_collection from chromy.handlers.delete_collection import ( handle_delete_collection, handle_delete_records, ) from chromy.handlers.import_data import handle_import from chromy.handlers.list_collections import handle_list_collections from chromy.handlers.query import handle_query CommandT = TypeVar("CommandT") class HandlerTests(unittest.TestCase): @staticmethod def _fixture_path(path: str) -> str: return str(Path(path).resolve()) def test_list_collections_prints_empty_message(self) -> None: with patch( "chromy.handlers.list_collections.list_collections", return_value=[] ): exit_code, output = _capture_output( handle_list_collections, ) self.assertEqual(exit_code, 0) self.assertEqual(output, "No collections found.\n") def test_list_collections_prints_collection_names(self) -> None: with patch( "chromy.handlers.list_collections.list_collections", return_value=["notes", "plays"], ): exit_code, output = _capture_output( handle_list_collections, ) self.assertEqual(exit_code, 0) self.assertEqual(output, "· notes\n· plays\n") def test_create_collection_uses_typed_input(self) -> None: with patch( "chromy.handlers.create_collection.create_collection", return_value="notes", ) as create_collection: exit_code, output = _capture_output( handle_create_collection, "notes", ) create_collection.assert_called_once_with("notes") self.assertEqual(exit_code, 0) self.assertEqual(output, "Created: collection 'notes'.\n") def test_delete_collection_uses_typed_input(self) -> None: with patch("chromy.handlers.delete_collection.delete_collection") as delete: exit_code, output = _capture_output( handle_delete_collection, "notes", ) delete.assert_called_once_with("notes") self.assertEqual(exit_code, 0) self.assertEqual(output, "Deleted collection 'notes'.\n") def test_count_collection_uses_typed_input(self) -> None: with patch( "chromy.handlers.count_collection.count_collection", return_value=7, ) as count: exit_code, output = _capture_output( handle_count_collection, "notes", ) count.assert_called_once_with("notes") self.assertEqual(exit_code, 0) self.assertEqual(output, "The 'notes' collection contains 7 records.\n") def test_import_data_uses_typed_input(self) -> None: with patch( "chromy.handlers.import_data.ingest_file", return_value=3, ) as ingest_file: exit_code, output = _capture_output( handle_import, "notes", ["romeo_and_juliet.txt"], ) ingest_file.assert_called_once_with( "notes", self._fixture_path("romeo_and_juliet.txt"), ) self.assertEqual(exit_code, 0) self.assertEqual( output, "Added 3 records from 'romeo_and_juliet.txt' to collection 'notes'.\n" "Imported 1 file(s) successfully; 0 failed.\n", ) def test_import_data_continues_after_missing_file(self) -> None: with patch( "chromy.handlers.import_data.ingest_file", return_value=3, ) as ingest_file: exit_code, output = _capture_output( handle_import, "notes", ["missing.txt", "romeo_and_juliet.txt"], ) ingest_file.assert_called_once_with( "notes", self._fixture_path("romeo_and_juliet.txt"), ) self.assertEqual(exit_code, 1) self.assertEqual( output, "Error: The file 'missing.txt' was not found.\n" "Added 3 records from 'romeo_and_juliet.txt' to collection 'notes'.\n" "Imported 1 file(s) successfully; 1 failed.\n", ) def test_import_data_rejects_non_text_files(self) -> None: with patch( "chromy.handlers.import_data.is_probably_text_file", return_value=False, ): exit_code, output = _capture_output( handle_import, "notes", ["romeo_and_juliet.txt"], ) self.assertEqual(exit_code, 1) self.assertEqual( output, "Error: The file 'romeo_and_juliet.txt' is not a text file.\n" "Imported 0 file(s) successfully; 1 failed.\n", ) def test_import_data_deduplicates_files(self) -> None: with patch( "chromy.handlers.import_data.ingest_file", return_value=3, ) as ingest_file: exit_code, output = _capture_output( handle_import, "notes", ["README.md", "./README.md"], ) ingest_file.assert_called_once_with( "notes", self._fixture_path("README.md"), ) self.assertEqual(exit_code, 0) self.assertEqual( output, "Added 3 records from 'README.md' to collection 'notes'.\n" "Imported 1 file(s) successfully; 0 failed.\n", ) def test_import_data_suppresses_per_file_output_with_progress(self) -> None: progress = MagicMock() progress.__enter__.return_value = progress progress.__exit__.return_value = None progress.console.print = print progress.add_task.return_value = 1 with ( patch("chromy.handlers.import_data.ingest_file", side_effect=[3, 2]), patch( "chromy.handlers.import_data._should_show_progress", return_value=True, ), patch("chromy.handlers.import_data.Progress", return_value=progress), ): exit_code, output = _capture_output( handle_import, "notes", ["romeo_and_juliet.txt", "README.md"], ) self.assertEqual(exit_code, 0) self.assertEqual(output, "Imported 2 file(s) successfully; 0 failed.\n") def test_import_data_truncates_long_file_names_in_progress(self) -> None: progress = MagicMock() progress.__enter__.return_value = progress progress.__exit__.return_value = None progress.console.print = print progress.add_task.return_value = 1 with ( patch( "chromy.handlers.import_data._get_absolute_path", side_effect=[ "/tmp/this_is_a_very_long_file_name.txt", self._fixture_path("README.md"), "/tmp/this_is_a_very_long_file_name.txt", self._fixture_path("README.md"), ], ), patch("chromy.handlers.import_data._import_one", return_value=3), patch( "chromy.handlers.import_data._should_show_progress", return_value=True, ), patch("chromy.handlers.import_data.Progress", return_value=progress), ): handle_import( "notes", ["this_is_a_very_long_file_name.txt", "README.md"], ) progress.update.assert_any_call( 1, description="Importing [bold]this_is_a_very_lo...[/]...", ) def test_query_uses_typed_input(self) -> None: query_result = {"ids": [["1"]], "documents": [["hello"]]} with ( patch("chromy.handlers.query.run_query", return_value=query_result) as run, patch( "chromy.handlers.query.format_query_result", return_value=["Query results:", "1"], ) as format_result, ): exit_code, output = _capture_output( handle_query, "notes", "hello", ) run.assert_called_once_with("notes", "hello") format_result.assert_called_once_with(query_result) self.assertEqual(exit_code, 0) self.assertEqual(output, "Query results:\n1\n") def test_delete_records_parses_where_filter(self) -> None: with patch( "chromy.handlers.delete_collection.delete_data", return_value=2, ) as delete_data: exit_code, output = _capture_output( handle_delete_records, "notes", " file_name = play.txt ", ) delete_data.assert_called_once_with("notes", {"file_name": "play.txt"}) self.assertEqual(exit_code, 0) self.assertEqual( output, "Deleted 2 record(s) from collection 'notes' where file_name=play.txt.\n", ) def test_delete_records_rejects_invalid_where_filter(self) -> None: with self.assertRaisesRegex( ValueError, "Invalid --where value. Expected =.", ): handle_delete_records("notes", "file_name") def _capture_output( handler: Callable[..., int], *arguments: object, ) -> tuple[int, str]: output = io.StringIO() with redirect_stdout(output): exit_code = handler(*arguments) return exit_code, output.getvalue() if __name__ == "__main__": unittest.main()