fixed test cases

Author: akashmangoai

adalflow/tests/test_lancedb_retriver.py

@@ -1,123 +1,141 @@
 import unittest
-from unittest.mock import Mock, MagicMock
+from unittest.mock import Mock, patch
 import numpy as np
+import pandas as pd
 from adalflow.components.retriever import LanceDBRetriever
 from adalflow.core.embedder import Embedder
-from adalflow.core.types import RetrieverOutput, Document
+from unittest import mock
+from adalflow.core.types import EmbedderOutput, RetrieverOutput
 
-# Mock LanceDB and PyArrow imports since they are specific to LanceDB
-lancedb = MagicMock()
-pa = MagicMock()
+# Helper function to create dummy embeddings
+def create_dummy_embeddings(num_embeddings, dim):
+    return np.random.rand(num_embeddings, dim).astype(np.float32)
 
 class TestLanceDBRetriever(unittest.TestCase):
     def setUp(self):
-        # Basic configuration
         self.dimensions = 128
+        self.top_k = 5
+        self.single_query = ["sample query"]
         self.embedder = Mock(spec=Embedder)
-        self.db_uri = "/tmp/test_lancedb"
 
-        # Mock embedding output with a simple structure
-        self.dummy_embeddings = np.random.rand(10, self.dimensions).astype(np.float32)
-        self.embedder.return_value.data = [
-            Mock(embedding=embedding) for embedding in self.dummy_embeddings
-        ]
-
-        # Initialize LanceDBRetriever with mocked embedder
-        self.retriever = LanceDBRetriever(
-            embedder=self.embedder, dimensions=self.dimensions, db_uri=self.db_uri
+        # Mock embedder to return dummy embeddings
+        self.dummy_embeddings = create_dummy_embeddings(10, self.dimensions)
+        self.embedder.return_value = EmbedderOutput(
+            data=[Mock(embedding=emb) for emb in self.dummy_embeddings[:len(self.single_query)]]
         )
 
-        # Mock LanceDB table and connection
-        self.retriever.db.create_table = MagicMock(return_value=Mock())
-        self.retriever.table = self.retriever.db.create_table.return_value
+        with patch("lancedb.connect") as mock_db_connect:
+            self.mock_db = mock_db_connect.return_value
+            self.mock_table = Mock()
+            self.mock_db.create_table.return_value = self.mock_table
+            self.retriever = LanceDBRetriever(
+                embedder=self.embedder,
+                dimensions=self.dimensions,
+                db_uri="/tmp/lancedb",
+                top_k=self.top_k
+            )
 
     def test_initialization(self):
-        # Check dimensions and embedder assignment
         self.assertEqual(self.retriever.dimensions, self.dimensions)
-        self.assertEqual(self.retriever.top_k, 5)
+        self.assertEqual(self.retriever.top_k, self.top_k)
+        self.mock_db.create_table.assert_called_once()
 
     def test_add_documents(self):
-        # Sample documents
-        documents = [{"content": f"Document {i}"} for i in range(5)]
+        documents = [{"content": f"Document {i}"} for i in range(10)]
+        embeddings = create_dummy_embeddings(len(documents), self.dimensions)
 
-        # Mock LanceDB table add method
-        self.retriever.table.add = MagicMock()
+        # Mock embedding output
+        self.embedder.return_value = EmbedderOutput(
+            data=[Mock(embedding=embedding) for embedding in embeddings]
+        )
 
-        # Add documents to LanceDBRetriever
         self.retriever.add_documents(documents)
+        self.assertEqual(self.mock_table.add.call_count, 1)
+        args, _ = self.mock_table.add.call_args
+        self.assertEqual(len(args[0]), len(documents))
+
+    def test_add_documents_no_documents(self):
+        self.retriever.add_documents([])
+        self.mock_table.add.assert_not_called()
 
-        # Ensure add method was called
-        self.retriever.table.add.assert_called_once()
-        # Verify embeddings were passed to LanceDB add method
-        added_data = self.retriever.table.add.call_args[0][0]
-        self.assertEqual(len(added_data), len(documents))
-        self.assertIn("vector", added_data[0])
-        self.assertIn("content", added_data[0])
-
-    def test_retrieve(self):
-        # Prepare a sample query and mocked search result from LanceDB
-        query = "test query"
-        dummy_scores = [0.9, 0.8, 0.7]
-        dummy_indices = [0, 1, 2]
-
-        # Set up mock search result as if it was retrieved from LanceDB
-        self.retriever.table.search = MagicMock(return_value=Mock())
-        self.retriever.table.search().limit().to_pandas.return_value = Mock(
-            index=dummy_indices, _distance=dummy_scores
+    def test_retrieve_single_query(self):
+        query = "sample query"
+        query_embedding = create_dummy_embeddings(1, self.dimensions)[0]
+
+        # Mock embedding for query
+        self.embedder.return_value = EmbedderOutput(
+            data=[Mock(embedding=query_embedding)]
         )
 
-        # Retrieve top-k results for the query
-        result = self.retriever.retrieve(query)
+        # Mock search results from LanceDB as pandas DataFrame
+        results_df = pd.DataFrame({
+            "index": [0, 1, 2],
+            "_distance": [0.1, 0.2, 0.3]
+        })
+        self.mock_table.search.return_value.limit.return_value.to_pandas.return_value = results_df
 
-        # Check if retrieve method returns expected output structure
-        self.assertIsInstance(result, list)
-        self.assertEqual(len(result), 1)
+        result = self.retriever.retrieve(query)
         self.assertIsInstance(result[0], RetrieverOutput)
-        self.assertEqual(result[0].query, query)
-        self.assertEqual(result[0].doc_indices, dummy_indices)
-        self.assertEqual(result[0].doc_scores, dummy_scores)
+        self.assertEqual(len(result[0].doc_indices), 3)
+        self.assertEqual(len(result[0].doc_scores), 3)
+        self.assertListEqual(result[0].doc_indices, [0, 1, 2])
+        self.assertListEqual(result[0].doc_scores, [0.1, 0.2, 0.3])
 
     def test_retrieve_multiple_queries(self):
-        # Prepare multiple queries and mocked search result
         queries = ["query 1", "query 2"]
-        dummy_scores = [[0.9, 0.8], [0.85, 0.75]]
-        dummy_indices = [[0, 1], [2, 3]]
-
-        # Set up mock for each query's result
-        self.retriever.table.search().limit().to_pandas.side_effect = [
-            Mock(index=dummy_indices[0], _distance=dummy_scores[0]),
-            Mock(index=dummy_indices[1], _distance=dummy_scores[1]),
-        ]
-
-        # Retrieve for multiple queries
-        results = self.retriever.retrieve(queries)
-
-        # Verify the structure and content of the results
-        self.assertEqual(len(results), len(queries))
-        for i, result in enumerate(results):
-            self.assertEqual(result.query, queries[i])
-            self.assertEqual(result.doc_indices, dummy_indices[i])
-            self.assertEqual(result.doc_scores, dummy_scores[i])
-
-    def test_empty_document_addition(self):
-        # Ensure warning log for empty document list
-        with self.assertLogs(level='WARNING'):
-            self.retriever.add_documents([])
+        query_embeddings = create_dummy_embeddings(len(queries), self.dimensions)
 
-    def test_retrieve_with_empty_query(self):
-        # Check empty query handling, expecting a list with empty RetrieverOutput
-        result = self.retriever.retrieve("")
-        self.assertEqual(result, [RetrieverOutput(doc_indices=[], doc_scores=[], query="")])
+        # Mock embedding for queries
+        self.embedder.return_value = EmbedderOutput(
+            data=[Mock(embedding=embedding) for embedding in query_embeddings]
+        )
+
+        # Mock search results for each query
+        results_df = pd.DataFrame({
+            "index": [0, 1, 2],
+            "_distance": [0.1, 0.2, 0.3]
+        })
+        self.mock_table.search.return_value.limit.return_value.to_pandas.return_value = results_df
 
-    def test_add_documents_embedding_failure(self):
-        # Simulate embedding failure
-        self.embedder.side_effect = Exception("Embedding failure")
-        documents = [{"content": "test document"}]
+        result = self.retriever.retrieve(queries)
+        self.assertEqual(len(result), len(queries))
+        for res in result:
+            self.assertIsInstance(res, RetrieverOutput)
+            self.assertEqual(len(res.doc_indices), 3)
+            self.assertEqual(len(res.doc_scores), 3)
 
-        with self.assertRaises(Exception) as context:
-            self.retriever.add_documents(documents)
+    def test_retrieve_with_empty_query(self):
+        # Mock the empty results DataFrame
+        self.mock_table.search.return_value.limit.return_value.to_pandas.return_value = pd.DataFrame({
+            "index": [],
+            "_distance": []
+        })
+
+    def test_retrieve_with_no_index(self):
+        empty_retriever = LanceDBRetriever(
+            embedder=self.embedder, dimensions=self.dimensions
+        )
+        with self.assertRaises(ValueError):
+            empty_retriever.retrieve("test query")
+
+    def test_overwrite_table_on_initialization(self):
+        with patch("lancedb.connect") as mock_db_connect:
+            mock_db = mock_db_connect.return_value
+            mock_table = Mock()
+            mock_db.create_table.return_value = mock_table
+
+            LanceDBRetriever(
+                embedder=self.embedder,
+                dimensions=self.dimensions,
+                db_uri="/tmp/lancedb",
+                overwrite=True
+            )
+            mock_db.create_table.assert_called_once_with(
+                "documents",
+                schema=mock.ANY,
+                mode="overwrite"
+            )
 
-        self.assertEqual(str(context.exception), "Embedding failure")
 
 if __name__ == "__main__":
     unittest.main()