diff --git a/tests/readme.md b/tests/readme.md
new file mode 100644
index 0000000..4ea6257
--- /dev/null
+++ b/tests/readme.md
@@ -0,0 +1,9 @@
+A testing plan including:
+
+- a plan for implementing concrete unit tests in each of ClassifAI's 4 modules - describing which componets to test,
+- table of key concerns for testing each module,
+- a corresponding first pass example of unit tests generated by Claude to test a specific class or feature detailed in the module unit testing plan.
+
+The above list describes content that we can use to implement sets of unit tests for each module, testing each independently of the other modules. Additional later tests could include end-to-end integration tests where we use multiple modules together (vectoriiser + VectorStore for example).
+
+Additionally, a `test_exports.py` test script within the parent folder. The purpose of this test is to ensure that each of the modules from the package that should be importable to a user are succsefully exported by the package. This ensure thats the the base API for the package is accessible and functioning.
\ No newline at end of file
diff --git a/tests/test_evaluation/readme.md b/tests/test_evaluation/readme.md
new file mode 100644
index 0000000..e6c2c76
--- /dev/null
+++ b/tests/test_evaluation/readme.md
@@ -0,0 +1,310 @@
+# Unit Testing Plan for Evaluation Module
+## Test Structure Overview
+1. parse_metrics Function Tests (test_parse_metrics.py)
+* Valid metric parsing:
+ * Single valid metric name returns dict with one entry
+ * Multiple valid metric names return dict with all entries
+ * Case insensitivity (accepts "accuracy", "ACCURACY", "Accuracy")
+ * Metric instances are correct type (e.g., ClassificationAccuracy for "accuracy")
+ * Dict keys match input metric names (lowercase)
+ * Dict values are Metric instances with evaluate() method
+* Invalid metric handling (ValueError):
+ * Single invalid metric name raises ValueError
+ * Invalid metric in list of valid ones raises error
+ * Error message includes the invalid metric name
+ * Error message includes list of valid metrics
+ * Empty string metric name raises error
+ * Whitespace-only metric name raises error
+ * Typos caught (e.g., "accuraccy" instead of "accuracy")
+* Edge cases:
+ * Empty list returns empty dict
+ * Duplicate metric names handled (e.g., ["accuracy", "accuracy"] returns dict with 1 entry)
+ * Mixed case like "MacRo_F1" handled correctly
+
+2. Evaluation Initialization Tests (test_evaluation_init.py)
+* Input validation (ground_truths DataFrame):
+ * DataFrame with correct schema (text, label columns) validates
+ * Missing 'text' column raises pandera SchemaError
+ * Missing 'label' column raises pandera SchemaError
+ * Wrong dtype for 'text' (not string) raises pandera error
+ * Wrong dtype for 'label' (not string) raises pandera error
+ * Type coercion works (int/float converted to string)
+ * Empty DataFrame validates (0 rows but correct schema)
+ * Extra columns allowed but ignored
+* Metrics validation (InvalidMetricError):
+ * Valid metric list ["accuracy", "macro_f1"] parses successfully
+ * Invalid metric name raises InvalidMetricError
+ * Empty metrics list accepted (no metrics to compute)
+ * InvalidMetricError has code "invalid_metric_error"
+ * InvalidMetricError context includes metrics list and cause
+* Batch size validation (DataValidationError):
+ * Default batch_size is correct
+ * Custom batch_size (e.g., 16, 32) stored correctly
+ * Negative batch_size raises error (if checked in init)
+ * Zero batch_size raises error (if checked in init)
+* save_output flag:
+ * Default is False
+ * Can be set to True
+ * Boolean type enforced
+* Attributes set correctly:
+ * self.ground_truths is a copy (not reference)
+ * self.ground_truths has new 'qid' column added (index as string)
+ * self.batch_size set correctly
+ * self.save_output set correctly
+ * self.metric_results initialized as empty dict
+ * self.parsed_metrics is dict of Metric instances
+* qid column generation:
+ * All rows get unique qid values
+ * qid values are strings
+ * qid values correspond to original index
+ * Original ground_truths data unchanged
+
+3. Evaluation.evaluate() Method Tests (test_evaluation_evaluate.py)
+* Input validation:
+ * vectorstores must be list (not tuple, dict)
+ * vectorstore_names must be list (not tuple, dict)
+ * Length of vectorstores must equal length of vectorstore_names
+ * Each item in vectorstores is VectorStore instance or callable
+ * Invalid VectorStore instance at index i caught with context
+ * All vectorstore_names must be strings
+ * Invalid name type at index i caught
+ * All vectorstore_names must be unique (no duplicates)
+ * output_file must be string or None
+ * output_file must end with ".csv" if provided
+ * overwrite must be boolean
+* File system handling (save_output=True):
+ * Default output_file is "evaluation_results.csv" if save_output=True and output_file=None
+ * Existing file raises error if overwrite=False
+ * Existing file overwritten if overwrite=True
+ * Parent directories created if don't exist
+ * Directory creation errors handled
+ * Results saved to correct file path
+ * CSV format is correct (columns: vectorstore_name, metric names)
+* File system handling (save_output=False):
+ * No file written even if output_file provided
+ * No error raised for existing files
+* VectorStore processing:
+ * Each vectorstore processed sequentially
+ * Callable vectorstores instantiated before use
+ * Callable instantiation errors wrapped in EvaluationError
+ * Instance vectorstores used directly
+ * Invalid callable (doesn't return VectorStore) caught
+ * VectorStore deleted from memory after use if callable
+* Search execution (_run_search):
+ * VectorStoreSearchInput created with qid and text columns
+ * vectorstore.search() called with correct params (n_results=1)
+ * batch_size from Evaluation passed to search
+ * Search failure wrapped in EvaluationError with context
+ * Search error context includes vectorstore_name
+* Results validation:
+ * _run_search returns DataFrame with SearchOutputSchema
+ * Required columns present (query_id, query_text, doc_label, doc_text, rank, score, ground_truth_label)
+ * ground_truth_label column merged correctly from ground_truths
+ * Schema validation enforces column types and constraints
+ * rank >= 0 enforced
+ * Pandera validation failures raised as SchemaError
+* Metric computation:
+ * Each parsed metric evaluated on results
+ * metric.evaluate() called with results DataFrame
+ * Metric results stored in self.metric_results
+ * Metric computation errors wrapped in EvaluationError
+ * Error context includes vectorstore_name and cause
+ * Metric results persist across vectorstores (accumulate)
+* Results aggregation:
+ * DataFrame created for each vectorstore with metric results
+ * Row indexed by vectorstore name
+ * All metrics included as columns
+ * Overall DataFrame concatenates results from all stores
+ * Row order matches vectorstore order
+ * No duplicate rows
+* Error handling and cleanup:
+ * Any step failure raises appropriate exception type
+ * Errors include context (vectorstore_name, cause)
+ * Processing continues until error (no partial results)
+ * Callable vectorstores cleaned up even on error (finally block)
+* Return value:
+ * Returns DataFrame with vectorstore names as index
+ * One row per vectorstore
+ * Columns are metric names
+ * Values are floats
+* Edge cases:
+ * Single vectorstore works
+ * Many vectorstores work (10+)
+ * Empty ground_truths (0 queries) handled
+ * Mixed callables and instances in same list
+
+4. Evaluation._run_search() Method Tests (test_evaluation_run_search.py)
+* Search input construction:
+ * VectorStoreSearchInput created with correct data
+ * 'id' column from self.ground_truths['qid']
+ * 'query' column from self.ground_truths['text']
+ * Input order matches ground_truths order
+* Search execution:
+ * vectorstore.search() called with SearchInput
+ * n_results=1 passed (top-1 search)
+ * batch_size from self.batch_size passed
+ * Search result is DataFrame
+* Merge operation:
+ * Results merged with ground_truths on query_id → qid
+ * Left join preserves all results
+ * Merge columns: qid, label from ground_truths
+ * 'label' column renamed to 'ground_truth_label'
+ * Column order correct in output
+* Output validation:
+ * Returns DataFrame with SearchOutputSchema
+ * All required columns present
+ * Pandera validation passed
+ * No extra/unexpected columns (or handled gracefully)
+ * Row count matches input queries
+* Error handling:
+ * VectorStore errors not caught (because they propagate to evaluate()
+ * Pandera validation errors propagate (SearchOutputSchema)
+ * Merge errors propagate (shouldn't happen with valid input)
+* Edge cases:
+ * Single query works
+ * Many queries work
+ * Query with no matching results handled (None/NaN for ground_truth_label)
+
+5. Metric Base Class Tests (test_metrics_base.py)
+* Metric ABC enforcement:
+ * Cannot instantiate Metric directly (abstract)
+ * Subclasses must implement evaluate()
+ * Subclass missing evaluate() raises TypeError
+* MetricResult dataclass:
+ * Can be instantiated with name and value
+ * Has repr that formats as "name: value"
+ * Value formatted to 4 decimal places
+
+6. ClassificationAccuracy Metric Tests (test_metrics_accuracy.py)
+* Correct predictions:
+ * All correct predictions → accuracy = 1.0
+ * No correct predictions → accuracy = 0.0
+ * 50% correct → accuracy = 0.5
+ * Accuracy = correct_count / total_count
+* Edge cases:
+ * Empty DataFrame (0 rows) → 0.0 (or error?)
+ * Single prediction correct → 1.0
+ * Single prediction wrong → 0.0
+ * NaN/None values handled
+ * Case sensitivity (doc_label vs ground_truth_label)
+* Output:
+ * Returns MetricResult
+ * name = "accuracy"
+ * value is float in [0.0, 1.0]
+
+7. ClassificationMacroRecall Metric Tests (test_metrics_macro_recall.py)
+* Single label:
+ * All true positives → recall = 1.0
+ * No true positives → recall = 0.0
+ * Recall = TP / (TP + FN) per label
+* Multiple labels:
+ * Recalls computed per label
+ * Macro recall is average of per-label recalls
+ * Recall for unseen label is 0.0
+ * Recall with zero denominator (TP=0, FN=0) is 0.0
+* Edge cases:
+ * Empty DataFrame → 0.0
+ * Single label → that label's recall
+ * Labels only in predictions (not ground truth) → FN=0, Recall=1.0
+ * Labels only in ground truth (not predictions) → TP=0, FN>0, Recall=0.0
+* Output:
+ * Returns MetricResult with name "macro_recall"
+ * Value is float in [0.0, 1.0]
+
+8. ClassificationMacroPrecision Metric Tests (test_metrics_macro_precision.py)
+* Single label:
+ * All true positives → precision = 1.0
+ * No true positives → precision = 0.0
+ * Precision = TP / (TP + FP) per label
+* Multiple labels:
+ * Precisions computed per label
+ * Macro precision is average of per-label precisions
+ * Precision for unseen label is 0.0
+ * Precision with zero denominator (TP=0, FP=0) is 0.0
+* Edge cases:
+ * Empty DataFrame → 0.0
+ * Single label → that label's precision
+ * False positives only → precision = 0.0
+* Output:
+ * Returns MetricResult with name "macro_precision"
+ * Value is float in [0.0, 1.0]
+
+9. ClassificationMacroF1 Metric Tests (test_metrics_macro_f1.py)
+* F1 calculation:
+ * F1 = 2 * (precision * recall) / (precision + recall)
+ * F1 with zero denominator = 0.0
+ * Perfect precision and recall → F1 = 1.0
+ * Zero precision and recall → F1 = 0.0
+* Multiple labels:
+ * F1 computed per label
+ * Macro F1 is average of per-label F1s
+ * F1 for unseen label is 0.0
+* Edge cases:
+ * Empty DataFrame → 0.0
+ * Single label → that label's F1
+ * Precision = 0, Recall > 0 → F1 = 0.0
+* Output:
+ * Returns MetricResult with name "macro_f1"
+ * Value is float in [0.0, 1.0]
+
+10. Schema Validation Tests (test_evaluation_schemas.py)
+* GroundTruthSchema:
+ * Accepts 'text' and 'label' columns (string type)
+ * Coerces types (int/float → string)
+ * Rejects missing columns
+ * Rejects wrong dtypes (without coercion possible)
+ * Extra columns allowed
+* SearchOutputSchema:
+ * All required columns present
+ * Correct dtypes
+ * rank >= 0 constraint enforced
+ * score is float
+ * Coercion applied
+ * Rejects missing columns
+ * Rejects invalid types
+
+11. Other Edge Cases and Error Handling (test_evaluation_edge_cases.py)
+* DataFrame edge cases:
+ * 0 queries (empty ground_truths)
+ * 1 query
+ * Many queries (1000+)
+ * Special characters in labels
+ * Very long text strings
+ * Unicode in labels/text
+ * Whitespace in labels
+* Metric edge cases:
+ * All predictions correct
+ * All predictions wrong
+ * Perfect imbalance (1 label dominates)
+ * Many labels (10+)
+ * Predictions never match ground truth
+ * Labels only exist in predictions
+ * Labels only exist in ground truth
+* Exception scenarios:
+ * VectorStore raises exception during search
+ * Callable raises exception during instantiation
+ * Metric raises exception during evaluate
+ * File save fails (permission error)
+ * Merge operation produces unexpected shape
+ * Pandera validation fails on search results
+
+
+
+## Key Testing Considerations
+
+| **Aspect** | **Strategy** |
+|--------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| VectorStore Mocking | Mock VectorStore.search() to return realistic SearchOutputSchema DataFrames; control results to test different accuracy/metric scenarios |
+| DataFrame Creation | Create realistic ground_truths and search result DataFrames with all required columns; test both happy path and edge cases |
+| Pandera Validation | Test schema acceptance/rejection; verify type coercion works; test constraint enforcement (rank >= 0) |
+| Callable Handling | Mock callables that return VectorStore instances; test callable errors separately from VectorStore errors; verify cleanup in finally block |
+| Metric Computation | Mock metrics to control results; test metric.evaluate() called with correct data; verify MetricResult objects correct |
+| File I/O | Mock filesystem operations to avoid creating real files in tests; test path logic (directory creation, CSV format) separately from core logic |
+| Integration Flow | Test full workflow with real (small) DataFrames and mocked VectorStore; verify data flows through each step correctly |
+| Error Propagation | Verify exceptions wrapped with correct context; test error messages include vectorstore_name; verify cleanup happens even on error |
+| Metric Formulas | Hand-calculate expected metric values for simple test cases; verify computed values match; test edge cases (zero denominators, NaN, etc.) |
+| Label Handling | Test with various label distributions (balanced, imbalanced, single label, many labels); test missing/NaN labels |
+| Case Sensitivity | Test that column names are matched correctly (doc_label vs ground_truth_label); test label values are case-sensitive |
+| Type Coercion | Verify Pandera coercion works (int→str, float→str); test when coercion would fail (e.g., complex objects) |
+| Performance | Don't test with huge datasets (1M rows); use small (10-100 row) test DataFrames; integration tests run quickly |
+
diff --git a/tests/test_evaluation/test_parse_metrics_example.py b/tests/test_evaluation/test_parse_metrics_example.py
new file mode 100644
index 0000000..27bc51d
--- /dev/null
+++ b/tests/test_evaluation/test_parse_metrics_example.py
@@ -0,0 +1,190 @@
+"""Unit tests for parse_metrics function."""
+
+from __future__ import annotations
+
+import pytest
+
+from classifai.evaluation.main import parse_metrics
+from classifai.evaluation.metrics import (
+ ClassificationAccuracy,
+ ClassificationMacroF1,
+ ClassificationMacroPrecision,
+ ClassificationMacroRecall,
+ Metric,
+)
+
+
+class TestParseMetricsValidInput:
+ """Tests for parse_metrics with valid inputs."""
+
+ def test_parse_single_valid_metric(self):
+ """Test parsing a single valid metric name."""
+ result = parse_metrics(["accuracy"])
+
+ assert isinstance(result, dict)
+ assert len(result) == 1
+ assert "accuracy" in result
+ assert isinstance(result["accuracy"], ClassificationAccuracy)
+
+ def test_parse_multiple_valid_metrics(self):
+ """Test parsing multiple valid metric names."""
+ result = parse_metrics(["accuracy", "macro_recall", "macro_precision", "macro_f1"])
+
+ NEXT_ANSWER = 4
+ assert len(result) == NEXT_ANSWER
+ assert "accuracy" in result
+ assert "macro_recall" in result
+ assert "macro_precision" in result
+ assert "macro_f1" in result
+ assert isinstance(result["accuracy"], ClassificationAccuracy)
+ assert isinstance(result["macro_recall"], ClassificationMacroRecall)
+ assert isinstance(result["macro_precision"], ClassificationMacroPrecision)
+ assert isinstance(result["macro_f1"], ClassificationMacroF1)
+
+ def test_parse_metrics_case_insensitive(self):
+ """Test that metric names are case insensitive."""
+ result_lower = parse_metrics(["accuracy"])
+ result_upper = parse_metrics(["ACCURACY"])
+ result_mixed = parse_metrics(["Accuracy"])
+
+ assert "accuracy" in result_lower
+ assert "ACCURACY" in result_upper
+ assert "Accuracy" in result_mixed
+ assert isinstance(result_lower["accuracy"], ClassificationAccuracy)
+ assert isinstance(result_upper["ACCURACY"], ClassificationAccuracy)
+ assert isinstance(result_mixed["Accuracy"], ClassificationAccuracy)
+
+ def test_parse_metrics_dict_keys_match_input_names(self):
+ """Test that dict keys match the input metric names exactly."""
+ result = parse_metrics(["ACCURACY", "Macro_Recall"])
+
+ # Keys should match the case of input
+ assert "ACCURACY" in result
+ assert "Macro_Recall" in result
+
+ def test_parse_metrics_returns_metric_instances(self):
+ """Test that returned values are Metric instances."""
+ result = parse_metrics(["accuracy", "macro_f1"])
+
+ for metric in result.values():
+ assert isinstance(metric, Metric)
+ assert hasattr(metric, "evaluate")
+ assert callable(metric.evaluate)
+
+ def test_parse_empty_list(self):
+ """Test parsing an empty metric list."""
+ result = parse_metrics([])
+
+ assert isinstance(result, dict)
+ assert len(result) == 0
+
+
+class TestParseMetricsInvalidInput:
+ """Tests for parse_metrics with invalid inputs."""
+
+ def test_parse_invalid_metric_name_raises_error(self):
+ """Test that invalid metric name raises ValueError."""
+ with pytest.raises(ValueError) as exc_info:
+ parse_metrics(["invalid_metric"])
+
+ error_msg = str(exc_info.value)
+ assert "invalid_metric" in error_msg.lower()
+ assert "Invalid metric" in error_msg
+
+ def test_parse_invalid_metric_error_includes_valid_metrics(self):
+ """Test that error message includes list of valid metrics."""
+ with pytest.raises(ValueError) as exc_info:
+ parse_metrics(["wrong_name"])
+
+ error_msg = str(exc_info.value)
+ assert "accuracy" in error_msg.lower()
+ assert "macro_recall" in error_msg.lower() or "recall" in error_msg.lower()
+ assert "macro_precision" in error_msg.lower() or "precision" in error_msg.lower()
+ assert "macro_f1" in error_msg.lower() or "f1" in error_msg.lower()
+
+ def test_parse_typo_in_metric_name_raises_error(self):
+ """Test that typos in metric names are caught."""
+ with pytest.raises(ValueError):
+ parse_metrics(["accuraccy"]) # typo
+
+ with pytest.raises(ValueError):
+ parse_metrics(["macro_recal"]) # typo
+
+ with pytest.raises(ValueError):
+ parse_metrics(["f1_macro"]) # wrong order
+
+ def test_parse_empty_string_metric_raises_error(self):
+ """Test that empty string metric name raises error."""
+ with pytest.raises(ValueError):
+ parse_metrics([""])
+
+ def test_parse_whitespace_only_metric_raises_error(self):
+ """Test that whitespace-only metric name raises error."""
+ with pytest.raises(ValueError):
+ parse_metrics([" "])
+
+ def test_parse_invalid_metric_in_list_with_valid_ones(self):
+ """Test that invalid metric in list with valid ones raises error."""
+ with pytest.raises(ValueError) as exc_info:
+ parse_metrics(["accuracy", "invalid_metric", "macro_f1"])
+
+ error_msg = str(exc_info.value)
+ assert "invalid_metric" in error_msg.lower()
+
+ def test_parse_invalid_metric_preserves_partial_results_in_error(self):
+ """Test that error includes context about the failing metric."""
+ with pytest.raises(ValueError) as exc_info:
+ parse_metrics(["accuracy", "bad_metric"])
+
+ # Error should reference the bad metric, even if accuracy was already parsed
+ assert "bad_metric" in str(exc_info.value).lower()
+
+
+class TestParseMetricsEdgeCases:
+ """Tests for edge cases in parse_metrics."""
+
+ def test_parse_duplicate_metric_names(self):
+ """Test handling of duplicate metric names."""
+ result = parse_metrics(["accuracy", "accuracy"])
+
+ # Should return dict with single entry (dict keys are unique)
+ # or allow duplicates? Implementation dependent
+ assert "accuracy" in result
+ assert isinstance(result["accuracy"], ClassificationAccuracy)
+
+ def test_parse_all_available_metrics(self):
+ """Test parsing all available metrics at once."""
+ all_metrics = ["accuracy", "macro_recall", "macro_precision", "macro_f1"]
+ expected_metric_count = len(all_metrics)
+ result = parse_metrics(all_metrics)
+
+ assert len(result) == expected_metric_count
+ for metric_name in all_metrics:
+ assert metric_name in result
+
+ def test_parse_mixed_case_metrics(self):
+ """Test parsing metrics with various case combinations."""
+ result = parse_metrics(["ACCURACY", "Macro_Recall", "mACRO_PRECISION", "macro_F1"])
+
+ NEXT_ANSWER = 4
+ assert len(result) == NEXT_ANSWER
+ assert "ACCURACY" in result
+ assert "Macro_Recall" in result
+ assert "mACRO_PRECISION" in result
+ assert "macro_F1" in result
+
+ def test_parse_metrics_with_underscores(self):
+ """Test that metric names with underscores are handled correctly."""
+ result = parse_metrics(["macro_recall", "macro_precision", "macro_f1"])
+
+ assert "macro_recall" in result
+ assert "macro_precision" in result
+ assert "macro_f1" in result
+
+ def test_parse_metrics_with_spaces_raises_error(self):
+ """Test that metric names with spaces are invalid."""
+ with pytest.raises(ValueError):
+ parse_metrics(["macro recall"]) # space in name
+
+ with pytest.raises(ValueError):
+ parse_metrics(["accuracy "]) # trailing space
diff --git a/tests/test_exports.py b/tests/test_exports.py
new file mode 100644
index 0000000..e3ba27e
--- /dev/null
+++ b/tests/test_exports.py
@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import importlib
+import warnings
+
+import classifai
+
+
+def test_package_version_is_exposed():
+ assert classifai.__version__
+
+
+def test_vectorisers_public_exports_are_importable():
+ module = importlib.import_module("classifai.vectorisers")
+
+ assert module.VectoriserBase
+ assert module.HuggingFaceVectoriser
+ assert module.GcpVectoriser
+ assert module.OllamaVectoriser
+
+
+def test_indexers_public_exports_are_importable():
+ module = importlib.import_module("classifai.indexers")
+
+ assert module.VectorStore
+ assert module.VectorStoreEmbedInput
+ assert module.VectorStoreEmbedOutput
+ assert module.VectorStoreSearchInput
+ assert module.VectorStoreSearchOutput
+ assert module.VectorStoreReverseSearchInput
+ assert module.VectorStoreReverseSearchOutput
+
+
+def test_servers_public_exports_are_importable():
+ module = importlib.import_module("classifai.servers")
+
+ assert module.get_router
+ assert module.get_server
+ assert module.make_endpoints
+ assert module.run_server
+
+
+def test_evaluation_exports_are_importable():
+ module = importlib.import_module("classifai.evaluation")
+
+ assert module.Evaluation
+
+
+def test_evaluation_import_emits_future_warning():
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ module = importlib.import_module("classifai.evaluation")
+ importlib.reload(module)
+
+ assert any(item.category is FutureWarning for item in caught)
diff --git a/tests/test_indexers/readme.md b/tests/test_indexers/readme.md
new file mode 100644
index 0000000..f428b89
--- /dev/null
+++ b/tests/test_indexers/readme.md
@@ -0,0 +1,183 @@
+# Unit Testing Plan for Indexers Module
+## Test Structure Overview
+1. Dataclass Tests (test_indexers_dataclasses.py)
+* VectorStoreSearchInput:
+ * Valid dict/DataFrame converts and validates correctly
+ * Schema validation enforces column types
+ * Missing required columns raises validation error
+ * Type coercion works (strings, etc.)
+ * Property accessors work (id, query) / return correct series
+ * Empty inputs handled correctly
+* VectorStoreSearchOutput:
+ * Valid construction from dict/DataFrame
+ * Schema validation enforces column types
+ * Missing required columns raises validation error
+ * Rank column must be non-negative
+ * Score column accepts floats
+ * Property accessors work / return correct series
+ * Column ordering is preserved (queries broadcast down consecutive rows)
+ * Empty inputs handled correctly
+* VectorStoreEmbedInput/Output:
+ * Valid dict/DataFrame construction and validation
+ * Type coercion for id/text
+ * Embedding column accepts numpy arrays
+ * Empty Inputs/Output handled correctly
+* VectorStoreReverseSearchInput/Output:
+ * Valid construction from dicts/DataFrames
+ * Empty Input/Output DataFrame handles correctly
+ * Schema validation works
+ * Property accessors function properly
+ * Missing required columns raises validation error
+
+2. VectorStore Initialization Tests (test_vectorstore_init.py)
+* Input validation (DataValidationError):
+ * file_name must be non-empty string
+ * data_type validation (only "csv" supported)
+ * vectoriser must be VectoriserBase instance
+ * batch_size must be positive integer
+ * meta_data must be dict or None
+ * hooks must be dict or None
+ * output_dir must be string or None
+* File system handling (ConfigurationError):
+ * Input file must exist
+ * Output directory creation works
+ * overwrite flag prevents accidental overwrites
+ * gs:// paths require gcsfs (helpful error message)
+ * Invalid fsspec paths raise ConfigurationError
+* Index building (IndexBuildError):
+ * CSV file reads correctly
+ * UUID generation works
+ * Batch processing of embeddings
+ * Vectoriser failures wrapped appropriately
+ * Embeddings count matches batch size
+ * Metadata serialization to JSON
+ * Parquet file writing
+* skip_save flag:
+ * When True, no files written to disk
+ * When False, metadata.json and vectors.parquet created
+ * warning logged when output_dir set but skip_save=True
+
+3. VectorStore Search Tests (test_vectorstore_search.py)
+* Input validation (DataValidationError):
+ * query must be VectorStoreSearchInput
+ * n_results must be int >= 1
+ * batch_size must be int >= 1 or None
+ * Empty query raises error
+ * Vector store not initialized raises ConfigurationError
+* Search operation:
+ * Single query processes correctly
+ * Multiple queries in batch
+ * Similarity scores computed (dot-product)
+ * Top n_results returned per query
+ * Results ranked by score (descending)
+ * Output shape matches expected (n_queries * n_results rows)
+ * Metadata columns included in output
+ * Query batching with custom batch_size works
+* Error handling (VectorisationError/ClassifaiError):
+ * Query embedding failure
+ * Vectoriser.transform() exceptions wrapped
+ * Error context includes vectoriser class, batch info
+* Hooks integration:
+ * search_preprocess hook called before search
+ * search_postprocess hook called after search
+ * Hook failures raise HookError
+ * Multiple hooks in list processed in order
+ * Single hook converted to list automatically
+
+4. VectorStore Reverse Search Tests (test_vectorstore_reverse_search.py)
+* Input validation (DataValidationError):
+ * query must be VectorStoreReverseSearchInput
+ * max_n_results must be int >= 1 or -1
+ * Empty query raises error
+* Reverse search operation:
+ * Exact label matching works (default)
+ * Partial matching (prefix) when enabled
+ * max_n_results limits results per query
+ * max_n_results=-1 returns all matches
+ * Results include metadata columns
+ * Empty result sets handled (returns empty DataFrame with correct schema)
+ * Sorting by id and label works
+* Error handling:
+ * Vectoriser-independent (no embeddings needed)
+ * DataFrame join failures wrapped
+ * Error context includes max_n_results, query count
+* Hooks integration:
+ * reverse_search_preprocess hook calls before reverse search
+ * reverse_search_postprocess hook calls after reverse search
+ * Same checks and error handling as search
+
+5. VectorStore Embed Tests (test_vectorstore_embed.py)
+* Input validation (DataValidationError):
+ * query must be VectorStoreEmbedInput
+ * Invalid input type raises error
+* Embedding operation:
+ * Single text embeds correctly
+ * Multiple texts process correctly
+ * Output includes id, text, and embedding
+ * Embeddings are numpy arrays
+ * Output shape matches input count
+ * Vectoriser.transform() called with correct texts
+* Error handling (VectorisationError/ClassifaiError):
+ * Vectoriser failures wrapped with context
+ * Error includes vectoriser class, text count
+* Hooks integration:
+ * embed_preprocess hook called before embedding
+ * embed_postprocess hook called after embedding
+ * Same checks and error handling as search
+
+6. VectorStore Metadata Tests (test_vectorstore_metadata.py)
+* Metadata serialization (_save_metadata):
+ * JSON file created at correct path
+ * Contains all required fields (vectoriser_class, vector_shape, num_vectors, batch_size, created_at, meta_data)
+ * Type information preserved (str types → string names)
+ * Valid JSON format
+ * fsspec paths work (gs://, etc.)
+* Metadata loading (from_filespace):
+ * Metadata file read and parsed correctly
+ * Required keys validated
+ * Type deserialization works
+ * Backwards compatibility with v1.0.0 (missing batch_size)
+ * Default batch_size used when missing
+ * Warning logged for missing batch_size
+
+7. VectorStore from_filespace Tests (test_vectorstore_from_filespace.py)
+* Input validation (DataValidationError):
+ * folder_path must be non-empty string
+ * folder_path must be existing directory
+ * batch_size override must be int >= 1 or None
+ * hooks must be dict or None
+* File loading (IndexBuildError):
+ * metadata.json exists and valid
+ * vectors.parquet exists and valid
+ * Required columns present in parquet
+ * Parquet not empty
+ * Metadata can be deserialized
+* Configuration validation (ConfigurationError):
+ * Vectoriser class name matches metadata
+ * vectoriser must have callable .transform() method / inherit from base class
+ * fsspec paths (gs://) work with gcsfs
+ * Helpful error message when gcsfs missing
+* Instance construction:
+ * Instance created without calling init
+ * All attributes set correctly
+ * batch_size override works
+ * metadata.meta_data deserialized and set
+ * Vectoriser instance attached
+ * hooks parameter applied
+ * quiet_mode applied
+ * Instance is functional (can search/embed/reverse_search)
+
+
+## Key Testing Considerations
+| **Aspect** | **Strategy** |
+|--------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------|
+| Vectoriser Mocking | Mock `VectoriserBase` to return predictable embeddings; separate vectoriser testing from vectorstore testing |
+| File System | Mock `fsspec` for local/remote paths; test real local paths in integration tests; `gs://` tests optional/skipped without `gcsfs` |
+| Dataclass Validation | Test both valid and invalid inputs; verify `pandera` schema enforcement; test type coercion |
+| Large Datasets | Use small synthetic CSVs (< 100 rows); mock large searches with synthetic embeddings to avoid slow tests |
+| Similarity Computation | Verify dot-product calculations; test edge cases (zero embeddings, identical embeddings, single query) |
+| Hook System | Mock hooks that modify input/output; test hook chains; verify error propagation; test that single hooks auto-convert to lists |
+| Error Context | Verify all exceptions include relevant context (vectoriser class, batch info, file paths) without exposing secrets |
+| Save/Load Cycle | Test round-trip (create → save → load); verify metadata preservation; test backwards compatibility with old metadata format |
+| Empty/Edge Cases | Empty query results, single document, single query, all identical embeddings, `max_n_results > available docs` |
+| Quiet Mode | Verify progress bars suppressed; verify logging levels adjusted; test both `True` and `False` paths |
diff --git a/tests/test_indexers/test_vectorstore_init_example.py b/tests/test_indexers/test_vectorstore_init_example.py
new file mode 100644
index 0000000..1ebc16e
--- /dev/null
+++ b/tests/test_indexers/test_vectorstore_init_example.py
@@ -0,0 +1,481 @@
+"""Unit tests for VectorStore initialization."""
+
+from __future__ import annotations
+
+import json
+import tempfile
+from pathlib import Path
+from unittest.mock import Mock, patch
+
+import numpy as np
+import pytest
+
+from classifai.exceptions import (
+ ConfigurationError,
+ DataValidationError,
+ IndexBuildError,
+)
+from classifai.indexers import VectorStore
+from classifai.vectorisers import VectoriserBase
+
+
+class TestVectorStoreInitInputValidation:
+ """Tests for VectorStore initialization input validation."""
+
+ @pytest.fixture
+ def mock_vectoriser(self):
+ """Create a mocked vectoriser."""
+ vectoriser = Mock(spec=VectoriserBase)
+ vectoriser.transform.return_value = np.random.rand(3, 768)
+ return vectoriser
+
+ @pytest.fixture
+ def temp_csv_file(self):
+ """Create a temporary CSV file for testing."""
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
+ f.write("text,label\n")
+ f.write("hello world,111\n")
+ f.write("goodbye world,112\n")
+ f.write("test data,113\n")
+ temp_path = f.name
+ yield temp_path
+ Path(temp_path).unlink()
+
+ def test_init_with_valid_inputs(self, mock_vectoriser, temp_csv_file):
+ """Test successful initialization with valid inputs."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ vectorstore = VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ data_type="csv",
+ output_dir=temp_dir,
+ skip_save=True,
+ )
+
+ assert vectorstore.file_name == temp_csv_file
+ assert vectorstore.vectoriser == mock_vectoriser
+ NEXT_ANSWER = 128
+ assert vectorstore.batch_size == NEXT_ANSWER # default
+ assert vectorstore.meta_data == {} # default
+
+ def test_init_file_name_must_be_non_empty_string(self, mock_vectoriser):
+ """Test that file_name must be non-empty string."""
+ with pytest.raises(DataValidationError) as exc_info:
+ VectorStore(file_name="", vectoriser=mock_vectoriser)
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+ assert "file_name" in error.message.lower()
+
+ def test_init_file_name_must_exist(self, mock_vectoriser):
+ """Test that input file must exist."""
+ with pytest.raises(ConfigurationError) as exc_info:
+ VectorStore(
+ file_name="/nonexistent/path/file.csv",
+ vectoriser=mock_vectoriser,
+ )
+
+ error = exc_info.value
+ assert error.code == "configuration_error"
+ assert "not found" in error.message.lower() or "exist" in error.message.lower()
+
+ def test_init_data_type_must_be_csv(self, mock_vectoriser, temp_csv_file):
+ """Test that only 'csv' data_type is supported."""
+ with pytest.raises(DataValidationError) as exc_info:
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ data_type="parquet",
+ skip_save=True,
+ )
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+ assert "data_type" in error.message.lower()
+
+ def test_init_vectoriser_must_be_vectoriser_base_instance(self, temp_csv_file):
+ """Test that vectoriser must be VectoriserBase instance."""
+ invalid_vectoriser = "not a vectoriser"
+
+ with pytest.raises(DataValidationError) as exc_info:
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=invalid_vectoriser,
+ skip_save=True,
+ )
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+ assert "vectoriser" in error.message.lower()
+
+ def test_init_batch_size_must_be_positive_int(self, mock_vectoriser, temp_csv_file):
+ """Test that batch_size must be positive integer."""
+ with pytest.raises(DataValidationError) as exc_info:
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ batch_size=0,
+ skip_save=True,
+ )
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+ assert "batch_size" in error.message.lower()
+
+ def test_init_batch_size_negative_raises_error(self, mock_vectoriser, temp_csv_file):
+ """Test that negative batch_size raises error."""
+ with pytest.raises(DataValidationError):
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ batch_size=-1,
+ skip_save=True,
+ )
+
+ def test_init_meta_data_must_be_dict_or_none(self, mock_vectoriser, temp_csv_file):
+ """Test that meta_data must be dict or None."""
+ with pytest.raises(DataValidationError) as exc_info:
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ meta_data="invalid",
+ skip_save=True,
+ )
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+
+ def test_init_hooks_must_be_dict_or_none(self, mock_vectoriser, temp_csv_file):
+ """Test that hooks must be dict or None."""
+ with pytest.raises(DataValidationError) as exc_info:
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ hooks="invalid",
+ skip_save=True,
+ )
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+
+ def test_init_output_dir_must_be_string_or_none(self, mock_vectoriser, temp_csv_file):
+ """Test that output_dir must be string or None."""
+ with pytest.raises(DataValidationError) as exc_info:
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ output_dir=123,
+ skip_save=True,
+ )
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+
+
+class TestVectorStoreInitFileSystem:
+ """Tests for VectorStore file system handling during initialization."""
+
+ @pytest.fixture
+ def mock_vectoriser(self):
+ """Create a mocked vectoriser."""
+ vectoriser = Mock(spec=VectoriserBase)
+ vectoriser.transform.return_value = np.random.rand(3, 768)
+ return vectoriser
+
+ @pytest.fixture
+ def temp_csv_file(self):
+ """Create a temporary CSV file for testing."""
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
+ f.write("id,text\n")
+ f.write("1,hello world\n")
+ f.write("2,goodbye world\n")
+ f.write("3,test data\n")
+ temp_path = f.name
+ yield temp_path
+ Path(temp_path).unlink()
+
+ def test_init_creates_output_directory_if_not_exists(self, mock_vectoriser, temp_csv_file):
+ """Test that output_dir is created if it doesn't exist."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ output_path = Path(temp_dir) / "subdir" / "vectorstore"
+ assert not output_path.exists()
+
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ output_dir=str(output_path),
+ skip_save=True,
+ )
+
+ # Directory creation may or may not happen if skip_save=True
+ # The important thing is no error is raised
+
+ def test_init_overwrite_false_prevents_overwrite(self, mock_vectoriser, temp_csv_file):
+ """Test that overwrite=False prevents overwriting existing index."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ output_dir = Path(temp_dir) / "index"
+ output_dir.mkdir()
+
+ # Create existing metadata file
+ metadata_file = output_dir / "metadata.json"
+ metadata_file.write_text(json.dumps({"existing": "data"}))
+
+ with pytest.raises(ConfigurationError) as exc_info:
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ output_dir=str(output_dir),
+ overwrite=False,
+ )
+
+ error = exc_info.value
+ assert "overwrite" in error.message.lower()
+
+ def test_init_overwrite_true_allows_overwrite(self, mock_vectoriser, temp_csv_file):
+ """Test that overwrite=True allows overwriting existing index."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ output_dir = Path(temp_dir) / "index"
+ output_dir.mkdir()
+
+ # Create existing metadata file
+ metadata_file = output_dir / "metadata.json"
+ metadata_file.write_text(json.dumps({"existing": "data"}))
+
+ vectorstore = VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ output_dir=str(output_dir),
+ overwrite=True,
+ skip_save=True,
+ )
+
+ assert vectorstore is not None
+
+ def test_init_gcsfs_path_requires_gcsfs_library(self, mock_vectoriser, temp_csv_file):
+ """Test that gs:// paths require gcsfs and provide helpful error."""
+ with pytest.raises(ConfigurationError) as exc_info:
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ output_dir="gs://bucket/path",
+ skip_save=True,
+ )
+
+ error = exc_info.value
+ assert "gcsfs" in error.message.lower() or "google" in error.message.lower()
+
+
+class TestVectorStoreInitIndexBuilding:
+ """Tests for VectorStore index building during initialization."""
+
+ @pytest.fixture
+ def mock_vectoriser(self):
+ """Create a mocked vectoriser that returns embeddings."""
+ vectoriser = Mock(spec=VectoriserBase)
+ # Return embeddings with shape (batch_size, 768)
+ vectoriser.transform.side_effect = lambda texts: np.random.rand(
+ len(texts) if isinstance(texts, list) else 1, 768
+ )
+ return vectoriser
+
+ @pytest.fixture
+ def temp_csv_file(self):
+ """Create a temporary CSV file for testing."""
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
+ f.write("id,text,label\n")
+ f.write("1,hello world,positive\n")
+ f.write("2,goodbye world,negative\n")
+ f.write("3,test data,neutral\n")
+ temp_path = f.name
+ yield temp_path
+ Path(temp_path).unlink()
+
+ def test_init_reads_csv_file_correctly(self, mock_vectoriser, temp_csv_file):
+ """Test that CSV file is read correctly during initialization."""
+ vectorstore = VectorStore( # noqa: F841
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ skip_save=True,
+ )
+
+ # Verify vectoriser was called with text column
+ assert mock_vectoriser.transform.called
+ # Should have been called at least once with the texts
+ call_args = mock_vectoriser.transform.call_args_list
+ assert len(call_args) > 0
+
+ def test_init_uuid_generation_for_each_row(self, mock_vectoriser, temp_csv_file):
+ """Test that UUID is generated for each CSV row."""
+ vectorstore = VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ skip_save=True,
+ )
+
+ # Check that internal data has unique IDs
+ assert hasattr(vectorstore, "_index_data") or hasattr(vectorstore, "index_data")
+ # All UUIDs should be unique
+ index_attr = getattr(vectorstore, "_index_data", None) or getattr(vectorstore, "index_data", None)
+ if index_attr is not None and hasattr(index_attr, "index"):
+ NEXT_ANSWER = 3
+ assert len(index_attr.index.unique()) == NEXT_ANSWER # 3 rows in CSV
+
+ def test_init_batch_processing_of_embeddings(self, mock_vectoriser, temp_csv_file):
+ """Test that embeddings are processed in batches."""
+ vectorstore = VectorStore( # noqa: F841
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ batch_size=2,
+ skip_save=True,
+ )
+
+ # Vectoriser should be called multiple times (batched)
+ # With 3 texts and batch_size=2, should be called at least 2 times
+ NEXT_ANSWER = 2
+ assert mock_vectoriser.transform.call_count >= NEXT_ANSWER
+
+ def test_init_vectoriser_failure_raises_index_build_error(self, temp_csv_file):
+ """Test that vectoriser failures are wrapped in IndexBuildError."""
+ mock_vectoriser = Mock(spec=VectoriserBase)
+ mock_vectoriser.transform.side_effect = Exception("Vectoriser failed")
+
+ with pytest.raises(IndexBuildError) as exc_info:
+ VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ skip_save=True,
+ )
+
+ error = exc_info.value
+ assert error.code == "index_build_error"
+ assert "vectoriser" in error.message.lower()
+
+ def test_init_embeddings_count_matches_batch_size(self, mock_vectoriser, temp_csv_file):
+ """Test that returned embeddings match requested batch size."""
+ vectorstore = VectorStore( # noqa: F841
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ batch_size=2,
+ skip_save=True,
+ )
+
+ # All calls should return correct number of embeddings
+ for call in mock_vectoriser.transform.call_args_list:
+ texts = call[0][0] # First positional arg
+ embeddings = mock_vectoriser.transform.return_value
+ assert embeddings.shape[0] == len(texts)
+
+ def test_init_metadata_serialization(self, mock_vectoriser, temp_csv_file):
+ """Test that metadata is properly serialized."""
+ meta_data = {"source": "test", "version": "1.0"}
+
+ vectorstore = VectorStore(
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ meta_data=meta_data,
+ skip_save=True,
+ )
+
+ # Verify metadata is stored
+ assert vectorstore.meta_data == meta_data
+
+ def test_init_parquet_file_writing(self, mock_vectoriser, temp_csv_file):
+ """Test that embeddings are written to parquet format."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ vectorstore = VectorStore( # noqa: F841
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ output_dir=temp_dir,
+ skip_save=False,
+ )
+
+ # Check that vectors.parquet exists
+ parquet_file = Path(temp_dir) / "vectors.parquet"
+ assert parquet_file.exists()
+
+ def test_init_metadata_json_file_writing(self, mock_vectoriser, temp_csv_file):
+ """Test that metadata is written to JSON file."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ vectorstore = VectorStore( # noqa: F841
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ output_dir=temp_dir,
+ skip_save=False,
+ )
+
+ # Check that metadata.json exists
+ metadata_file = Path(temp_dir) / "metadata.json"
+ assert metadata_file.exists()
+
+ # Verify it's valid JSON
+ metadata = json.loads(metadata_file.read_text())
+ assert "vectoriser_class" in metadata
+ assert "vector_shape" in metadata
+
+
+class TestVectorStoreInitSkipSaveFlag:
+ """Tests for VectorStore skip_save flag behavior."""
+
+ @pytest.fixture
+ def mock_vectoriser(self):
+ """Create a mocked vectoriser."""
+ vectoriser = Mock(spec=VectoriserBase)
+ vectoriser.transform.return_value = np.random.rand(3, 768)
+ return vectoriser
+
+ @pytest.fixture
+ def temp_csv_file(self):
+ """Create a temporary CSV file for testing."""
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
+ f.write("id,text\n")
+ f.write("1,hello world\n")
+ f.write("2,goodbye world\n")
+ f.write("3,test data\n")
+ temp_path = f.name
+ yield temp_path
+ Path(temp_path).unlink()
+
+ def test_init_skip_save_true_no_files_written(self, mock_vectoriser, temp_csv_file):
+ """Test that no files are written when skip_save=True."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ vectorstore = VectorStore( # noqa: F841
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ output_dir=temp_dir,
+ skip_save=True,
+ )
+
+ # Check that no files were created
+ files = list(Path(temp_dir).glob("*"))
+ assert len(files) == 0
+
+ def test_init_skip_save_false_writes_files(self, mock_vectoriser, temp_csv_file):
+ """Test that files are written when skip_save=False."""
+ with tempfile.TemporaryDirectory() as temp_dir:
+ vectorstore = VectorStore( # noqa: F841
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ output_dir=temp_dir,
+ skip_save=False,
+ )
+
+ # Check that files were created
+ parquet_file = Path(temp_dir) / "vectors.parquet"
+ metadata_file = Path(temp_dir) / "metadata.json"
+ assert parquet_file.exists()
+ assert metadata_file.exists()
+
+ def test_init_skip_save_true_with_output_dir_logs_warning(self, mock_vectoriser, temp_csv_file):
+ """Test that warning is logged when output_dir set but skip_save=True."""
+ with tempfile.TemporaryDirectory() as temp_dir, patch("classifai.indexers.main.logger") as mock_logger:
+ vectorstore = VectorStore( # noqa: F841
+ file_name=temp_csv_file,
+ vectoriser=mock_vectoriser,
+ output_dir=temp_dir,
+ skip_save=True,
+ )
+
+ # Verify warning was logged
+ mock_logger.warning.assert_called()
+ warning_msg = mock_logger.warning.call_args[0][0]
+ assert "skip_save" in warning_msg.lower() or "not saved" in warning_msg.lower()
diff --git a/tests/test_servers/readme.md b/tests/test_servers/readme.md
new file mode 100644
index 0000000..ebfb7d9
--- /dev/null
+++ b/tests/test_servers/readme.md
@@ -0,0 +1,204 @@
+# Unit Testing Plan for Servers Module
+## Test Structure Overview
+1. Router Creation Tests (test_get_router.py)
+* Input validation (DataValidationError):
+ * vector_stores must be a list (not tuple, dict, etc.)
+ * endpoint_names must be a list (not tuple, dict, etc.)
+ * Length of vector_stores must match length of endpoint_names
+ * All endpoint_names must be non-empty strings
+ * No whitespace-only strings allowed
+ * endpoint_names must be unique (no duplicates)
+ * Empty lists raise appropriate error
+* VectorStore validation (ConfigurationError):
+ * Each item in vector_stores must be VectorStore instance
+ * Invalid type at specific index raises error with context
+ * Mixed valid/invalid stores caught at first invalid
+ * Error context includes the invalid index and type
+* Router creation:
+ * Router is successfully created and returned
+ * Router is FastAPI APIRouter instance
+ * Endpoints are registered for each vector store
+ * Correct number of sub-routers created
+ * Docs endpoint "/" redirects to "/docs"
+ * Router has correct tags for each endpoint
+* Edge cases:
+ * Single vector store works
+ * Many vector stores work (10+)
+ * Special characters in endpoint names handled
+ * Case sensitivity in endpoint names preserved
+
+2. Server Creation Tests (test_get_server.py)
+* Input validation (delegates to get_router):
+ * Same validation as get_router tested indirectly
+ * Invalid inputs raise same errors
+* FastAPI app creation:
+ * FastAPI instance returned
+ * App title set correctly ("ClassifAI API Server")
+ * App description set correctly
+ * App version matches __version__
+ * OpenAPI tags created for each endpoint name
+ * Each tag has correct name and description
+ * Router included in app
+ * Docs and redoc endpoints available
+* Integration:
+ * Router endpoints accessible through app
+ * All three endpoint types (search, embed, reverse_search) present
+ * Correct number of paths registered
+ * Tags properly organized in OpenAPI spec
+
+3. Server Runtime Tests (test_run_server.py)
+* Input validation (DataValidationError):
+ * port must be integer
+ * port must be >= 1
+ * port must be <= 65535
+ * Negative port raises error
+ * Port 0 raises error
+ * Port 65536 raises error
+ * host_ip must be string
+ * Empty host_ip validation
+* Log level validation (DataValidationError):
+ * Valid log levels: "debug", "info", "warning", "error", "critical"
+ * Invalid log level raises error
+ * Case sensitivity (lowercase required)
+ * Error message includes valid options
+ * Empty string raises error
+* Server startup (mocked uvicorn):
+ * uvicorn.run() called with correct parameters
+ * Port passed correctly
+ * Host IP passed correctly
+ * Log level passed correctly
+ * App created before uvicorn.run()
+ * Correct number of calls to uvicorn.run()
+* demo_mode flag:
+ * When False, app title remains "ClassifAI API Server"
+ * When True, app title changes to "ClassifAI API Demo Server"
+ * When True, app description changes to demo description
+ * _set_demo_defaults() called only when True
+ * Other app settings unaffected by demo_mode
+* Error handling:
+ * Invalid port caught before uvicorn.run()
+ * Invalid log_level caught before uvicorn.run()
+ * Validation errors propagate correctly
+
+4. Endpoint Creation Tests (test_endpoint_creation.py)
+* Search endpoint creation (_create_search_endpoint):
+ * Endpoint registered at /{name}/search
+ * HTTP method is POST
+ * Endpoint summary includes endpoint name
+ * Endpoint description includes endpoint name
+ * n_results query parameter has correct constraints (ge=1)
+ * n_results default value is 10
+ * Endpoint callable and returns SearchResponseBody
+* Embed endpoint creation (_create_embedding_endpoint):
+ * Endpoint registered at /{name}/embed
+ * HTTP method is POST
+ * Endpoint summary includes endpoint name
+ * Endpoint description includes endpoint name
+ * No query parameters
+ * Endpoint callable and returns EmbedResponseBody
+* Reverse search endpoint creation (_create_reverse_search_endpoint):
+ * Endpoint registered at /{name}/reverse_search
+ * HTTP method is POST
+ * Endpoint summary includes endpoint name
+ * Endpoint description includes endpoint name
+ * max_n_results query parameter with correct constraints
+ * max_n_results can be -1 (return all) or >= 1
+ * max_n_results default value is 100
+ * partial_match query parameter is boolean
+ * partial_match default is False
+ * Manual validation for max_n_results < 1 (when != -1) raises HTTPException(422)
+ * Endpoint callable and returns ReverseSearchResponseBody
+
+5. Endpoint Functional Tests (test_endpoint_functionality.py)
+* Search endpoint functionality:
+ * Extracts ids and queries from request
+ * Creates VectorStoreSearchInput with correct data
+ * Calls vectorstore.search() with correct params
+ * Calls convert_search_dataframe_to_pydantic_response()
+ * Returns formatted result as JSON
+ * Vectorstore search failure propagates as 500
+ * Invalid input format raises 422
+ * Empty queries handled
+* Embed endpoint functionality:
+ * Extracts ids and texts from request
+ * Creates VectorStoreEmbedInput with correct data
+ * Calls vectorstore.embed() with correct params
+ * Calls convert_embedding_dataframe_to_pydantic_response()
+ * Returns formatted result as JSON
+ * Vectorstore embed failure propagates
+ * Invalid input format raises 422
+ * Empty texts handled
+* Reverse search endpoint functionality:
+ * Extracts ids and doc_labels from request
+ * Creates VectorStoreReverseSearchInput with correct data
+ * Calls vectorstore.reverse_search() with correct params
+ * Calls convert_reverse_search_dataframe_to_pydantic_response()
+ * Returns formatted result as JSON
+ * max_n_results validation (< 1 when != -1) raises HTTPException(422)
+ * Vectorstore reverse_search failure propagates
+ * Invalid input format raises 422
+
+6. Response Conversion Tests (test_response_conversions.py)
+* Search response conversion (convert_search_dataframe_to_pydantic_response):
+ * Valid DataFrame converts to SearchResponseBody
+ * Grouped by query_id correctly
+ * Each group becomes SearchResponseSet
+ * Required columns present in output (query_id, query_text, entries)
+ * Metadata columns extracted and included dynamically
+ * Hook columns identified and included dynamically
+ * Rank column included (0-indexed or 1-indexed depending on implementation)
+ * Score column included as float
+ * Empty groups handled
+ * Multiple queries grouped separately
+ * Metadata dict respected (only columns in meta_data included)
+* Reverse search response conversion (convert_reverse_search_dataframe_to_pydantic_response):
+ * Valid DataFrame converts to ReverseSearchResponseBody
+ * Includes original_input to ensure all inputs in response
+ * Grouped by input id
+ * Empty result sets for inputs with no matches (still included in response)
+ * Each group becomes ReverseSearchResponseSet
+ * Required columns present (input_id, searched_doc_label, entries)
+ * Metadata columns extracted and included dynamically
+ * Hook columns identified and included dynamically
+ * searched_doc_label taken from first row of group
+ * doc_label and doc_text included in entries
+ * Multiple inputs handled separately
+* Embed response conversion (convert_embedding_dataframe_to_pydantic_response):
+ * Valid DataFrame converts to EmbedResponseBody
+ * Each row becomes EmbedResponseEntry
+ * id column included
+ * text column included
+ * embedding column converted to list (numpy array → list)
+ * Hook columns identified and included dynamically
+ * No meta_data parameter required
+ * Multiple embeddings handled
+ * Empty DataFrame returns empty data list
+ * Embedding dtype preserved (floats)
+
+7. Make Endpoints Tests (test_make_endpoints.py)
+* Router/app routing (make_endpoints):
+ * Accepts APIRouter or FastAPI app
+ * Creates sub_router for each vector store
+ * Sub_router has correct prefix (/{name})
+ * Sub_router has correct tags ([name])
+ * All three endpoint types created for each store
+ * Sub_routers included in main router/app
+ * Correct number of total endpoints
+ * Endpoints accessible at correct paths
+
+## Key Testing Considerations
+| Aspect | Strategy |
+|-----------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| VectorStore Mocking | Mock VectorStore instances with search/embed/reverse_search methods returning realistic DataFrames; avoid actual vectorstore initialization |
+| FastAPI Testing | Use TestClient from starlette to make HTTP requests; test actual endpoint behavior, not just function calls |
+| Pydantic Models | Verify request models parse correctly; test both valid and invalid JSON payloads; verify response models serialize correctly |
+| Input Validation | Test all validation branches in get_router/get_server/run_server; verify context information in errors |
+| Query Parameters | Test ge/le constraints on query parameters; test default values; test invalid type coercion |
+| DataFrame Conversions | Use real pandas DataFrames with all required columns; test edge cases (empty, single row, many rows); verify column ordering preserved |
+| Metadata Handling | Mock meta_data dict; verify only specified columns included; test missing metadata columns gracefully ignored |
+| Hook Columns | Test DataFrames with extra columns; verify they're included in responses; test without hook columns |
+| HTTPException Handling| Mock vectorstore to raise exceptions; verify they propagate correctly; test manual validation (e.g., max_n_results check) |
+| URL Construction | Verify endpoint paths are correct (/{name}/search, etc.); test special characters in endpoint names; test path collision prevention |
+| Logging | Mock logger; verify appropriate logs at startup (Starting ClassifAI Router, Generating ClassifAI API, Registering endpoints) |
+| Demo Mode | Test both True/False branches; verify only title/description changed, not functionality |
+
diff --git a/tests/test_servers/test_get_router_example.py b/tests/test_servers/test_get_router_example.py
new file mode 100644
index 0000000..73d2541
--- /dev/null
+++ b/tests/test_servers/test_get_router_example.py
@@ -0,0 +1,361 @@
+"""Unit tests for get_router function."""
+
+from __future__ import annotations
+
+from unittest.mock import Mock, patch
+
+import pytest
+from fastapi import APIRouter
+
+from classifai.exceptions import ConfigurationError, DataValidationError
+from classifai.indexers import VectorStore
+from classifai.servers import get_router
+
+
+class TestGetRouterInputValidation:
+ """Tests for get_router input validation."""
+
+ @pytest.fixture
+ def mock_vectorstore(self):
+ """Create a mocked VectorStore."""
+ return Mock(spec=VectorStore)
+
+ def test_get_router_valid_inputs(self, mock_vectorstore):
+ """Test successful router creation with valid inputs."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = ["search_endpoint"]
+
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names)
+
+ assert isinstance(router, APIRouter)
+ assert router is not None
+
+ def test_get_router_vector_stores_must_be_list(self, mock_vectorstore):
+ """Test that vector_stores must be a list."""
+ vector_stores = (mock_vectorstore,) # tuple, not list
+ endpoint_names = ["search_endpoint"]
+
+ with pytest.raises(DataValidationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+ assert "list" in error.message.lower()
+ assert error.context["vector_stores_type"] == "tuple"
+
+ def test_get_router_endpoint_names_must_be_list(self, mock_vectorstore):
+ """Test that endpoint_names must be a list."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = ("search_endpoint",) # tuple, not list
+
+ with pytest.raises(DataValidationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+ assert "list" in error.message.lower()
+ assert error.context["endpoint_names_type"] == "tuple"
+
+ def test_get_router_vector_stores_dict_not_list(self, mock_vectorstore):
+ """Test that vector_stores cannot be a dict."""
+ vector_stores = {"store": mock_vectorstore}
+ endpoint_names = ["search_endpoint"]
+
+ with pytest.raises(DataValidationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+
+ def test_get_router_lengths_must_match(self, mock_vectorstore):
+ """Test that vector_stores and endpoint_names must have same length."""
+ vector_stores = [mock_vectorstore, mock_vectorstore]
+ endpoint_names = ["search_endpoint"] # Only 1 name for 2 stores
+
+ with pytest.raises(DataValidationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+ assert "match" in error.message.lower()
+ NEXT_ANSWER = 2
+ assert error.context["n_vector_stores"] == NEXT_ANSWER
+ NEXT_ANSWER = 1
+ assert error.context["n_endpoint_names"] == NEXT_ANSWER
+
+ def test_get_router_endpoint_names_must_be_non_empty_strings(self, mock_vectorstore):
+ """Test that all endpoint_names must be non-empty strings."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = [""] # Empty string
+
+ with pytest.raises(DataValidationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+ assert "non-empty" in error.message.lower() or "empty" in error.message.lower()
+
+ def test_get_router_endpoint_names_no_whitespace_only(self, mock_vectorstore):
+ """Test that endpoint_names cannot be whitespace-only."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = [" "] # Whitespace only
+
+ with pytest.raises(DataValidationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+
+ def test_get_router_endpoint_names_must_be_strings(self, mock_vectorstore):
+ """Test that all endpoint_names must be strings."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = [123] # Integer, not string
+
+ with pytest.raises(DataValidationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+
+ def test_get_router_endpoint_names_must_be_unique(self, mock_vectorstore):
+ """Test that endpoint_names must be unique."""
+ vector_stores = [mock_vectorstore, mock_vectorstore]
+ endpoint_names = ["search", "search"] # Duplicate names
+
+ with pytest.raises(DataValidationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.code == "validation_error"
+ assert "unique" in error.message.lower()
+
+ def test_get_router_empty_lists(self):
+ """Test that empty lists are handled."""
+ vector_stores = []
+ endpoint_names = []
+
+ # Empty lists are valid - no endpoints to create
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names)
+
+ assert isinstance(router, APIRouter)
+
+
+class TestGetRouterVectorStoreValidation:
+ """Tests for VectorStore instance validation in get_router."""
+
+ @pytest.fixture
+ def mock_vectorstore(self):
+ """Create a mocked VectorStore."""
+ return Mock(spec=VectorStore)
+
+ def test_get_router_each_store_must_be_vectorstore_instance(self, mock_vectorstore):
+ """Test that each item must be a VectorStore instance."""
+ vector_stores = [mock_vectorstore, "not a vectorstore"]
+ endpoint_names = ["store1", "store2"]
+
+ with pytest.raises(ConfigurationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.code == "configuration_error"
+ assert "VectorStore" in error.message
+ assert error.context["index"] == 1
+ assert error.context["vector_store_type"] == "str"
+
+ def test_get_router_invalid_store_at_first_position(self):
+ """Test that invalid store at first position raises error."""
+ vector_stores = [None]
+ endpoint_names = ["store1"]
+
+ with pytest.raises(ConfigurationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.context["index"] == 0
+
+ def test_get_router_invalid_store_at_middle_position(self, mock_vectorstore):
+ """Test that invalid store in middle of list is caught."""
+ vector_stores = [mock_vectorstore, 123, mock_vectorstore]
+ endpoint_names = ["store1", "store2", "store3"]
+
+ with pytest.raises(ConfigurationError) as exc_info:
+ get_router(vector_stores, endpoint_names)
+
+ error = exc_info.value
+ assert error.context["index"] == 1
+ assert error.context["vector_store_type"] == "int"
+
+ def test_get_router_multiple_stores_all_valid(self, mock_vectorstore):
+ """Test that multiple valid stores work correctly."""
+ mock_store1 = Mock(spec=VectorStore)
+ mock_store2 = Mock(spec=VectorStore)
+ mock_store3 = Mock(spec=VectorStore)
+
+ vector_stores = [mock_store1, mock_store2, mock_store3]
+ endpoint_names = ["store1", "store2", "store3"]
+
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names)
+
+ assert isinstance(router, APIRouter)
+
+
+class TestGetRouterRouterCreation:
+ """Tests for router creation and endpoint registration."""
+
+ @pytest.fixture
+ def mock_vectorstore(self):
+ """Create a mocked VectorStore."""
+ return Mock(spec=VectorStore)
+
+ def test_get_router_returns_apirouter(self, mock_vectorstore):
+ """Test that get_router returns an APIRouter instance."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = ["search"]
+
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names)
+
+ assert isinstance(router, APIRouter)
+
+ def test_get_router_docs_endpoint_exists(self, mock_vectorstore):
+ """Test that the docs endpoint "/" is registered."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = ["search"]
+
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names)
+
+ # Check that "/" endpoint is registered
+ routes = [route.path for route in router.routes]
+ assert "/" in routes
+
+ def test_get_router_docs_endpoint_redirects_to_docs(self, mock_vectorstore):
+ """Test that docs endpoint redirects to /docs."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = ["search"]
+
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names)
+
+ # Find the docs endpoint and verify it redirects to /docs
+ docs_route = None
+ for route in router.routes:
+ if route.path == "/":
+ docs_route = route
+ break
+
+ assert docs_route is not None
+ assert "GET" in docs_route.methods or "get" in str(docs_route)
+
+ def test_get_router_make_endpoints_called(self, mock_vectorstore):
+ """Test that make_endpoints is called with correct arguments."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = ["search"]
+
+ with patch("classifai.servers.main.make_endpoints") as mock_make:
+ router = get_router(vector_stores, endpoint_names)
+
+ # Verify make_endpoints was called with router and dict mapping
+ mock_make.assert_called_once()
+ call_args = mock_make.call_args
+ assert call_args[0][0] == router # First arg is router
+ assert isinstance(call_args[0][1], dict) # Second arg is dict
+ assert "search" in call_args[0][1] # Dict has endpoint name
+
+ def test_get_router_logging_info_called(self, mock_vectorstore):
+ """Test that logging.info is called on router creation."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = ["search"]
+
+ with patch("classifai.servers.main.logging.info") as mock_log:
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names) # noqa: F841
+
+ # Verify logging was called
+ assert mock_log.called
+ # Find the "Starting ClassifAI Router" log
+ log_messages = [call[0][0] for call in mock_log.call_args_list]
+ assert any("Router" in msg for msg in log_messages)
+
+ def test_get_router_special_characters_in_names(self, mock_vectorstore):
+ """Test that special characters in endpoint names are preserved."""
+ vector_stores = [mock_vectorstore, mock_vectorstore]
+ endpoint_names = ["search_v1", "search-v2"]
+
+ with patch("classifai.servers.main.make_endpoints") as mock_make:
+ router = get_router(vector_stores, endpoint_names) # noqa: F841
+
+ call_args = mock_make.call_args
+ stores_dict = call_args[0][1]
+ assert "search_v1" in stores_dict
+ assert "search-v2" in stores_dict
+
+
+class TestGetRouterEdgeCases:
+ """Tests for edge cases in get_router."""
+
+ @pytest.fixture
+ def mock_vectorstore(self):
+ """Create a mocked VectorStore."""
+ return Mock(spec=VectorStore)
+
+ def test_get_router_single_store_single_endpoint(self, mock_vectorstore):
+ """Test with single store and single endpoint."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = ["only_store"]
+
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names)
+
+ assert isinstance(router, APIRouter)
+
+ def test_get_router_many_stores(self, mock_vectorstore):
+ """Test with many stores and endpoints."""
+ mock_stores = [Mock(spec=VectorStore) for _ in range(10)]
+ endpoint_names = [f"store_{i}" for i in range(10)]
+
+ with patch("classifai.servers.main.make_endpoints") as mock_make:
+ router = get_router(mock_stores, endpoint_names) # noqa: F841
+
+ # Verify all stores and names are included
+ call_args = mock_make.call_args
+ stores_dict = call_args[0][1]
+ NEXT_ANSWER = 10
+ assert len(stores_dict) == NEXT_ANSWER
+
+ def test_get_router_case_sensitive_names(self, mock_vectorstore):
+ """Test that endpoint names are case sensitive."""
+ mock_store1 = Mock(spec=VectorStore)
+ mock_store2 = Mock(spec=VectorStore)
+
+ vector_stores = [mock_store1, mock_store2]
+ endpoint_names = ["Store", "store"] # Different cases
+
+ # Should succeed - they're different names
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names)
+
+ assert isinstance(router, APIRouter)
+
+ def test_get_router_endpoint_with_unicode_characters(self, mock_vectorstore):
+ """Test endpoint names with unicode characters."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = ["café"] # Unicode character
+
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names)
+
+ assert isinstance(router, APIRouter)
+
+ def test_get_router_long_endpoint_names(self, mock_vectorstore):
+ """Test with very long endpoint names."""
+ vector_stores = [mock_vectorstore]
+ endpoint_names = ["a" * 100] # Very long name
+
+ with patch("classifai.servers.main.make_endpoints"):
+ router = get_router(vector_stores, endpoint_names)
+
+ assert isinstance(router, APIRouter)
diff --git a/tests/test_vectorisers/readme.md b/tests/test_vectorisers/readme.md
new file mode 100644
index 0000000..ca091ee
--- /dev/null
+++ b/tests/test_vectorisers/readme.md
@@ -0,0 +1,69 @@
+# Unit Testing Plan for Vectorisers Module
+## Test Structure Overview
+
+1. Base Class Tests (test_vectoriser_base.py)
+* Verify VectoriserBase is abstract and cannot be instantiated
+* Verify transform method is abstract
+* Test that subclasses must implement transform
+
+
+2. HuggingFaceVectoriser Tests(test_huggingface_vectoriser.py)
+* Initialisation:
+ * Missing dependencies raise appropriate errors
+ * Valid model loads successfully
+ * Invalid model name raises ExternalServiceError
+ * Device selection (CPU/GPU) works correctly
+ * Bad device selection raises ConfigurationError
+ * trust_remote_code defaults to False
+ * Custom kwargs are passed through
+* Transform method:
+ * Single string input converts to list and processes
+ * List of strings processes correctly
+ * Returns 2D numpy array
+ * Output shape matches input count
+ * Tokenisation failures raise VectorisationError
+ * Model inference failures raise VectorisationError
+ * Pooling failures raise VectorisationError
+
+
+3. GcpVectoriser Tests (test_gcp_vectoriser.py)
+* Initialization:
+ * Missing dependencies raise appropriate errors
+ * project_id + location authentication works
+ * api_key authentication works
+ * Missing both auth methods raises ConfigurationError
+ * Providing both auth methods raises ConfigurationError
+ * Client initialisation failures raise ConfigurationError
+* Transform method:
+ * Single string input converts to list
+ * List processes correctly
+ * Returns 2D numpy array
+ * Output shape matches input count
+ * API request failures raise ExternalServiceError
+ * Unexpected response format raises VectorisationError
+
+
+4. OllamaVectoriser Tests (test_ollama_vectoriser.py)
+* Initialization:
+ * Missing dependencies raise appropriate errors
+ * Model name is stored correctly
+* Transform method:
+ * Single string input converts to list
+ * List processes correctly
+ * Returns 2D numpy array
+ * Service failures raise ExternalServiceError
+ * Response parsing failures raise VectorisationError
+
+
+
+## Key Testing Considerations
+
+
+| **Aspect** | **Strategy** |
+|------------------------|-----------------------------------------------------------------------------|
+| External Dependencies | Use `pytest-mock` or `unittest.mock` to patch external libraries (torch, transformers, ollama, google.genai) |
+| GPU/Device Testing | Mock `torch.cuda` to test both CPU and GPU branches if we are concerned with GPU compatibility |
+| API Responses | Mock service responses with realistic embedding data |
+| Error Cases | Test each exception path in try-except blocks |
+| Input Validation | Test both string and list inputs |
+| Output Validation | Verify numpy array shape, dtype, and content |
diff --git a/tests/test_vectorisers/test_huggingface_vectoriser_example.py b/tests/test_vectorisers/test_huggingface_vectoriser_example.py
new file mode 100644
index 0000000..7e95330
--- /dev/null
+++ b/tests/test_vectorisers/test_huggingface_vectoriser_example.py
@@ -0,0 +1,340 @@
+"""Unit tests for HuggingFaceVectoriser."""
+
+from __future__ import annotations
+
+from unittest.mock import Mock, patch
+
+import numpy as np
+import pytest
+
+from classifai.exceptions import ConfigurationError, ExternalServiceError, VectorisationError
+from classifai.vectorisers import HuggingFaceVectoriser
+
+
+class TestHuggingFaceVectoriserInitialization:
+ """Tests for HuggingFaceVectoriser initialization."""
+
+ @patch("classifai.vectorisers.huggingface.check_deps")
+ @patch("classifai.vectorisers.huggingface.AutoTokenizer")
+ @patch("classifai.vectorisers.huggingface.AutoModel")
+ def test_init_with_valid_model(self, mock_model, mock_tokenizer, mock_check_deps):
+ """Test successful initialization with a valid model name."""
+ mock_tokenizer_instance = Mock()
+ mock_model_instance = Mock()
+ mock_tokenizer.from_pretrained.return_value = mock_tokenizer_instance
+ mock_model.from_pretrained.return_value = mock_model_instance
+
+ vectoriser = HuggingFaceVectoriser("bert-base-uncased")
+
+ assert vectoriser.model_name == "bert-base-uncased"
+ assert vectoriser.tokenizer == mock_tokenizer_instance
+ assert vectoriser.model == mock_model_instance
+ mock_check_deps.assert_called_once_with(["transformers", "torch"], extra="huggingface")
+
+ @patch("classifai.vectorisers.huggingface.check_deps")
+ def test_init_missing_dependencies(self, mock_check_deps):
+ """Test initialization fails when required dependencies are missing."""
+ mock_check_deps.side_effect = ImportError("Missing dependency")
+
+ with pytest.raises(ImportError):
+ HuggingFaceVectoriser("bert-base-uncased")
+
+ @patch("classifai.vectorisers.huggingface.check_deps")
+ @patch("classifai.vectorisers.huggingface.AutoTokenizer")
+ @patch("classifai.vectorisers.huggingface.AutoModel")
+ def test_init_invalid_model_name_raises_external_service_error(self, mock_model, mock_tokenizer, mock_check_deps):
+ """Test that invalid model name raises ExternalServiceError."""
+ mock_tokenizer.from_pretrained.side_effect = Exception("Model not found")
+
+ with pytest.raises(ExternalServiceError) as exc_info:
+ HuggingFaceVectoriser("invalid-model-xyz")
+
+ error = exc_info.value
+ assert error.code == "external_service_error"
+ assert "Failed to load HuggingFace model/tokenizer" in error.message
+ assert error.context["vectoriser"] == "huggingface"
+ assert error.context["model"] == "invalid-model-xyz"
+
+ @patch("classifai.vectorisers.huggingface.check_deps")
+ @patch("classifai.vectorisers.huggingface.AutoTokenizer")
+ @patch("classifai.vectorisers.huggingface.AutoModel")
+ def test_init_with_custom_tokenizer_kwargs(self, mock_model, mock_tokenizer, mock_check_deps):
+ """Test initialization with custom tokenizer kwargs."""
+ mock_tokenizer_instance = Mock()
+ mock_model_instance = Mock()
+ mock_tokenizer.from_pretrained.return_value = mock_tokenizer_instance
+ mock_model.from_pretrained.return_value = mock_model_instance
+
+ custom_kwargs = {"trust_remote_code": True, "cache_dir": "/custom/path"}
+ HuggingFaceVectoriser("bert-base-uncased", tokenizer_kwargs=custom_kwargs)
+
+ # Verify trust_remote_code was preserved (not overridden to False)
+ call_kwargs = mock_tokenizer.from_pretrained.call_args[1]
+ assert call_kwargs["trust_remote_code"] is True
+ assert call_kwargs["cache_dir"] == "/custom/path"
+
+ @patch("classifai.vectorisers.huggingface.check_deps")
+ @patch("classifai.vectorisers.huggingface.AutoTokenizer")
+ @patch("classifai.vectorisers.huggingface.AutoModel")
+ def test_init_trust_remote_code_defaults_to_false(self, mock_model, mock_tokenizer, mock_check_deps):
+ """Test that trust_remote_code defaults to False for security."""
+ mock_tokenizer_instance = Mock()
+ mock_model_instance = Mock()
+ mock_tokenizer.from_pretrained.return_value = mock_tokenizer_instance
+ mock_model.from_pretrained.return_value = mock_model_instance
+
+ HuggingFaceVectoriser("bert-base-uncased")
+
+ tokenizer_call_kwargs = mock_tokenizer.from_pretrained.call_args[1]
+ model_call_kwargs = mock_model.from_pretrained.call_args[1]
+ assert tokenizer_call_kwargs["trust_remote_code"] is False
+ assert model_call_kwargs["trust_remote_code"] is False
+
+ @patch("classifai.vectorisers.huggingface.check_deps")
+ @patch("classifai.vectorisers.huggingface.AutoTokenizer")
+ @patch("classifai.vectorisers.huggingface.AutoModel")
+ def test_init_with_custom_model_revision(self, mock_model, mock_tokenizer, mock_check_deps):
+ """Test initialization with custom model revision."""
+ mock_tokenizer_instance = Mock()
+ mock_model_instance = Mock()
+ mock_tokenizer.from_pretrained.return_value = mock_tokenizer_instance
+ mock_model.from_pretrained.return_value = mock_model_instance
+
+ HuggingFaceVectoriser("bert-base-uncased", model_revision="dev")
+
+ tokenizer_call_kwargs = mock_tokenizer.from_pretrained.call_args[1]
+ model_call_kwargs = mock_model.from_pretrained.call_args[1]
+ assert tokenizer_call_kwargs["revision"] == "dev"
+ assert model_call_kwargs["revision"] == "dev"
+
+ @patch("classifai.vectorisers.huggingface.check_deps")
+ @patch("classifai.vectorisers.huggingface.AutoTokenizer")
+ @patch("classifai.vectorisers.huggingface.AutoModel")
+ def test_init_device_selection_with_explicit_device(self, mock_model, mock_tokenizer, mock_check_deps):
+ """Test initialization with explicit device selection."""
+ mock_tokenizer_instance = Mock()
+ mock_model_instance = Mock()
+ mock_tokenizer.from_pretrained.return_value = mock_tokenizer_instance
+ mock_model.from_pretrained.return_value = mock_model_instance
+
+ with patch("classifai.vectorisers.huggingface.torch") as mock_torch: # noqa: F841
+ mock_device = Mock()
+ vectoriser = HuggingFaceVectoriser("bert-base-uncased", device=mock_device)
+
+ assert vectoriser.device == mock_device
+ mock_model_instance.to.assert_called_once_with(mock_device)
+ mock_model_instance.eval.assert_called_once()
+
+ @patch("classifai.vectorisers.huggingface.check_deps")
+ @patch("classifai.vectorisers.huggingface.AutoTokenizer")
+ @patch("classifai.vectorisers.huggingface.AutoModel")
+ def test_init_device_selection_auto_defaults_to_gpu_if_available(self, mock_model, mock_tokenizer, mock_check_deps):
+ """Test that device auto-selection chooses GPU if available."""
+ mock_tokenizer_instance = Mock()
+ mock_model_instance = Mock()
+ mock_tokenizer.from_pretrained.return_value = mock_tokenizer_instance
+ mock_model.from_pretrained.return_value = mock_model_instance
+
+ with patch("classifai.vectorisers.huggingface.torch") as mock_torch:
+ mock_gpu_device = Mock()
+ mock_torch.cuda.is_available.return_value = True
+ mock_torch.device.return_value = mock_gpu_device
+
+ vectoriser = HuggingFaceVectoriser("bert-base-uncased", device=None)
+
+ mock_torch.cuda.is_available.assert_called_once()
+ mock_torch.device.assert_called_with("cuda")
+ assert vectoriser.device == mock_gpu_device
+
+ @patch("classifai.vectorisers.huggingface.check_deps")
+ @patch("classifai.vectorisers.huggingface.AutoTokenizer")
+ @patch("classifai.vectorisers.huggingface.AutoModel")
+ def test_init_device_selection_fallback_to_cpu(self, mock_model, mock_tokenizer, mock_check_deps):
+ """Test that device selection falls back to CPU when GPU unavailable."""
+ mock_tokenizer_instance = Mock()
+ mock_model_instance = Mock()
+ mock_tokenizer.from_pretrained.return_value = mock_tokenizer_instance
+ mock_model.from_pretrained.return_value = mock_model_instance
+
+ with patch("classifai.vectorisers.huggingface.torch") as mock_torch:
+ mock_cpu_device = Mock()
+ mock_torch.cuda.is_available.return_value = False
+ mock_torch.device.return_value = mock_cpu_device
+
+ vectoriser = HuggingFaceVectoriser("bert-base-uncased", device=None)
+
+ mock_torch.device.assert_called_with("cpu")
+ assert vectoriser.device == mock_cpu_device
+
+ @patch("classifai.vectorisers.huggingface.check_deps")
+ @patch("classifai.vectorisers.huggingface.AutoTokenizer")
+ @patch("classifai.vectorisers.huggingface.AutoModel")
+ def test_init_device_initialization_failure_raises_configuration_error(
+ self, mock_model, mock_tokenizer, mock_check_deps
+ ):
+ """Test that device initialization failure raises ConfigurationError."""
+ mock_tokenizer_instance = Mock()
+ mock_model_instance = Mock()
+ mock_tokenizer.from_pretrained.return_value = mock_tokenizer_instance
+ mock_model.from_pretrained.return_value = mock_model_instance
+ mock_model_instance.to.side_effect = RuntimeError("Device not available")
+
+ with patch("classifai.vectorisers.huggingface.torch") as mock_torch: # noqa: F841
+ with pytest.raises(ConfigurationError) as exc_info:
+ HuggingFaceVectoriser("bert-base-uncased", device=Mock())
+
+ error = exc_info.value
+ assert error.code == "configuration_error"
+ assert "Failed to initialise model on device" in error.message
+ assert error.context["vectoriser"] == "huggingface"
+
+
+class TestHuggingFaceVectoriserTransform:
+ """Tests for HuggingFaceVectoriser transform method."""
+
+ @pytest.fixture
+ def mock_vectoriser(self):
+ """Create a mocked HuggingFaceVectoriser instance."""
+ with (
+ patch("classifai.vectorisers.huggingface.check_deps"),
+ patch("classifai.vectorisers.huggingface.AutoTokenizer"),
+ patch("classifai.vectorisers.huggingface.AutoModel"),
+ ):
+ vectoriser = HuggingFaceVectoriser("bert-base-uncased")
+ vectoriser.tokenizer = Mock()
+ vectoriser.model = Mock()
+ vectoriser.device = Mock()
+ return vectoriser
+
+ def test_transform_single_string_converts_to_list(self, mock_vectoriser):
+ """Test that a single string input is converted to a list."""
+ mock_inputs = Mock()
+ mock_inputs.to.return_value = mock_inputs
+ mock_vectoriser.tokenizer.return_value = mock_inputs
+
+ mock_output = Mock()
+ mock_output.last_hidden_state = Mock()
+ mock_vectoriser.model.return_value = mock_output
+
+ with patch("classifai.vectorisers.huggingface.torch.nn.functional") as mock_f:
+ mock_f.normalize.return_value = np.array([[0.1, 0.2, 0.3]])
+ mock_vectoriser.transform("hello world")
+
+ call_args = mock_vectoriser.tokenizer.call_args[0][0]
+ assert isinstance(call_args, list)
+ assert call_args == ["hello world"]
+
+ def test_transform_list_of_strings_processes_correctly(self, mock_vectoriser):
+ """Test that a list of strings is processed correctly."""
+ texts = ["text1", "text2", "text3"]
+ mock_inputs = Mock()
+ mock_inputs.to.return_value = mock_inputs
+ mock_vectoriser.tokenizer.return_value = mock_inputs
+
+ mock_output = Mock()
+ mock_output.last_hidden_state = Mock()
+ mock_vectoriser.model.return_value = mock_output
+
+ with patch("classifai.vectorisers.huggingface.torch.nn.functional") as mock_f:
+ mock_f.normalize.return_value = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]])
+ result = mock_vectoriser.transform(texts)
+
+ call_args = mock_vectoriser.tokenizer.call_args[0][0]
+ assert call_args == texts
+ assert isinstance(result, np.ndarray)
+
+ def test_transform_returns_2d_numpy_array(self, mock_vectoriser):
+ """Test that transform returns a 2D numpy array."""
+ mock_inputs = Mock()
+ mock_inputs.to.return_value = mock_inputs
+ mock_vectoriser.tokenizer.return_value = mock_inputs
+
+ mock_output = Mock()
+ mock_output.last_hidden_state = Mock()
+ mock_vectoriser.model.return_value = mock_output
+
+ embeddings = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
+ with patch("classifai.vectorisers.huggingface.torch.nn.functional") as mock_f:
+ mock_f.normalize.return_value = embeddings
+ result = mock_vectoriser.transform(["text1", "text2"])
+
+ assert isinstance(result, np.ndarray)
+ NEXT_ANSWER = 2
+ assert result.ndim == NEXT_ANSWER
+ assert result.shape == (2, 3)
+
+ def test_transform_output_shape_matches_input_count(self, mock_vectoriser):
+ """Test that output shape matches the number of input texts."""
+ texts = ["a", "b", "c", "d", "e"]
+ embedding_dim = 768
+
+ mock_inputs = Mock()
+ mock_inputs.to.return_value = mock_inputs
+ mock_vectoriser.tokenizer.return_value = mock_inputs
+
+ mock_output = Mock()
+ mock_output.last_hidden_state = Mock()
+ mock_vectoriser.model.return_value = mock_output
+
+ embeddings = np.random.rand(len(texts), embedding_dim)
+ with patch("classifai.vectorisers.huggingface.torch.nn.functional") as mock_f:
+ mock_f.normalize.return_value = embeddings
+ result = mock_vectoriser.transform(texts)
+
+ assert result.shape[0] == len(texts)
+ assert result.shape[1] == embedding_dim
+
+ def test_transform_tokenization_failure_raises_vectorisation_error(self, mock_vectoriser):
+ """Test that tokenization failures raise VectorisationError."""
+ mock_vectoriser.tokenizer.side_effect = Exception("Tokenization failed")
+
+ with pytest.raises(VectorisationError) as exc_info:
+ mock_vectoriser.transform("some text")
+
+ error = exc_info.value
+ assert error.code == "vectorisation_error"
+ assert "Tokenization" in error.message or "tokenization" in str(error)
+
+ def test_transform_model_inference_failure_raises_vectorisation_error(self, mock_vectoriser):
+ """Test that model inference failures raise VectorisationError."""
+ mock_inputs = Mock()
+ mock_inputs.to.return_value = mock_inputs
+ mock_vectoriser.tokenizer.return_value = mock_inputs
+ mock_vectoriser.model.side_effect = Exception("Model inference failed")
+
+ with pytest.raises(VectorisationError) as exc_info:
+ mock_vectoriser.transform("some text")
+
+ error = exc_info.value
+ assert error.code == "vectorisation_error"
+
+ def test_transform_pooling_failure_raises_vectorisation_error(self, mock_vectoriser):
+ """Test that pooling/normalization failures raise VectorisationError."""
+ mock_inputs = Mock()
+ mock_inputs.to.return_value = mock_inputs
+ mock_vectoriser.tokenizer.return_value = mock_inputs
+
+ mock_output = Mock()
+ mock_output.last_hidden_state = Mock()
+ mock_vectoriser.model.return_value = mock_output
+
+ with patch("classifai.vectorisers.huggingface.torch.nn.functional") as mock_f:
+ mock_f.normalize.side_effect = Exception("Pooling failed")
+
+ with pytest.raises(VectorisationError) as exc_info:
+ mock_vectoriser.transform("some text")
+
+ error = exc_info.value
+ assert error.code == "vectorisation_error"
+
+ def test_transform_error_context_includes_model_info(self, mock_vectoriser):
+ """Test that error context includes model and vectoriser information."""
+ mock_vectoriser.model_name = "distilbert-base-uncased"
+ mock_vectoriser.tokenizer.side_effect = Exception("Tokenization failed")
+
+ with pytest.raises(VectorisationError) as exc_info:
+ mock_vectoriser.transform("some text")
+
+ error = exc_info.value
+ assert error.context["vectoriser"] == "huggingface"
+ assert error.context["model"] == "distilbert-base-uncased"