Addressing feedback and fixing tests

Fiona-Waters · Fiona-Waters · commit 04fd031c4088 · 2025-06-17T12:37:52.000+01:00
Signed-off-by: Fiona Waters &lt;fiwaters6@gmail.com&gt;
diff --git a/sdk/python/feast/rag_retriever.py b/sdk/python/feast/rag_retriever.py
@@ -33,22 +33,6 @@ def __init__(self):
         """Initialize the Feast index."""
         pass
 
-    def get_top_docs(self, query_vectors: np.ndarray, n_docs: int = 5):
-        """Get top documents (not implemented).
-
-        This method is required by the RagRetriever interface but is not used
-        as we override the retrieve method in FeastRAGRetriever.
-        """
-        raise NotImplementedError("get_top_docs is not yet implemented.")
-
-    def get_doc_dicts(self, doc_ids: List[str]):
-        """Get document dictionaries (not implemented).
-
-        This method is required by the RagRetriever interface but is not used
-        as we override the retrieve method in FeastRAGRetriever.
-        """
-        raise NotImplementedError("get_doc_dicts is not yet implemented.")
-
 
 class FeastRAGRetriever(RagRetriever):
     """RAG retriever implementation that uses Feast as a backend."""
@@ -114,12 +98,6 @@ def __init__(
             **kwargs,
         )
 
-        # if torch.cuda.is_available():
-        #     self.question_encoder.to(torch.device("cuda"))
-        #     self.generator_model.to(torch.device("cuda"))
-        # self.question_encoder = question_encoder
-        # self.generator_model = generator_model
-        # self.generator_tokenizer = generator_tokenizer
         self.feast_repo_path = feast_repo_path
         self.search_type = search_type.lower()
         self.format_document = format_document or self._default_format_document
diff --git a/sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py b/sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py
@@ -198,7 +198,6 @@ def test_dynamodb_online_store_update(repo_config, dynamodb_online_store):
     assert existing_tables is not None
     assert len(existing_tables) == 1
     assert existing_tables[0] == f"test_aws.{db_table_keep_name}"
-
     assert _get_tags(dynamodb_client, existing_tables[0]) == [
         {"Key": "some", "Value": "tag"}
     ]
diff --git a/sdk/python/tests/unit/test_rag_retriever.py b/sdk/python/tests/unit/test_rag_retriever.py
@@ -1,5 +1,5 @@
 import os
-from unittest.mock import MagicMock, Mock
+from unittest.mock import MagicMock, Mock, patch
 
 import numpy as np
 import pytest
@@ -18,7 +18,7 @@
 from tests.utils.rag_test_utils import MockVectorStore
 
 
-@pytest.fixture(autouse=True)
+@pytest.fixture
 def cleanup_milvus():
     """Ensure Milvus resources are cleaned up between tests."""
     yield
@@ -104,7 +104,7 @@ def forward(self, **kwargs):
 
 
 @pytest.fixture
-def feature_store(temp_dir):
+def simple_feature_store(temp_dir):
     """Create a simple feature store without auth for RAG tests."""
     return FeatureStore(
         config=RepoConfig(
@@ -117,7 +117,7 @@ def feature_store(temp_dir):
 
 
 @pytest.fixture
-def rag_retriever(feature_store):
+def rag_retriever(simple_feature_store):
     """Create a RAG retriever instance for testing."""
     # Import the required objects
     from tests.example_repos.example_feature_repo_1 import (
@@ -137,7 +137,7 @@ def rag_retriever(feature_store):
         question_encoder=MockModel(),
         generator_tokenizer=MockTokenizer(),
         generator_model=MockModel(),
-        feast_repo_path=str(feature_store.repo_path),
+        feast_repo_path=str(simple_feature_store.repo_path),
         feature_view=document_embeddings,
         features=[
             "document_embeddings:content",
@@ -153,7 +153,7 @@ def rag_retriever(feature_store):
 
     # Replace the vector store with our mock
     retriever._vector_store = MockVectorStore(
-        repo_path=str(feature_store.repo_path),
+        repo_path=str(simple_feature_store.repo_path),
         rag_view=document_embeddings,
         features=[
             "document_embeddings:content",
@@ -181,13 +181,13 @@ def test_feast_index_initialization():
     assert index is not None
 
 
-def test_feast_index_methods():
-    """Test FeastIndex methods raise NotImplementedError."""
-    index = FeastIndex()
-    with pytest.raises(NotImplementedError):
-        index.get_top_docs(np.array([1, 2, 3]))
-    with pytest.raises(NotImplementedError):
-        index.get_doc_dicts(["doc1", "doc2"])
+# def test_feast_index_methods():
+#     """Test FeastIndex methods raise NotImplementedError."""
+#     index = FeastIndex()
+#     with pytest.raises(NotImplementedError):
+#         index.get_top_docs(np.array([1, 2, 3]))
+#     with pytest.raises(NotImplementedError):
+#         index.get_doc_dicts(["doc1", "doc2"])
 
 
 def test_rag_retriever_initialization(rag_retriever):
@@ -200,7 +200,7 @@ def test_rag_retriever_initialization(rag_retriever):
     assert rag_retriever.format_document is not None  # Should have default formatter
 
 
-def test_rag_retriever_custom_format_document(feature_store):
+def test_rag_retriever_custom_format_document(simple_feature_store):
     """Test RAG retriever initialization with custom document formatter."""
     from tests.example_repos.example_feature_repo_1 import document_embeddings
 
@@ -212,7 +212,7 @@ def custom_formatter(doc):
         question_encoder=MockModel(),
         generator_tokenizer=MockTokenizer(),
         generator_model=MockModel(),
-        feast_repo_path=str(feature_store.repo_path),
+        feast_repo_path=str(simple_feature_store.repo_path),
         feature_view=document_embeddings,
         features=[
             "document_embeddings:content",
@@ -247,7 +247,7 @@ def test_default_format_document(rag_retriever):
     assert "Embeddings" not in formatted  # Vector should be skipped
 
 
-def test_rag_retriever_invalid_search_type(feature_store):
+def test_rag_retriever_invalid_search_type(simple_feature_store):
     """Test RAG retriever initialization with invalid search type."""
     from tests.example_repos.example_feature_repo_1 import (
         document_embeddings,
@@ -259,7 +259,7 @@ def test_rag_retriever_invalid_search_type(feature_store):
             question_encoder=MockModel(),
             generator_tokenizer=MockTokenizer(),
             generator_model=MockModel(),
-            feast_repo_path=str(feature_store.repo_path),
+            feast_repo_path=str(simple_feature_store.repo_path),
             feature_view=document_embeddings,
             features=["content", "title", "Embeddings"],
             search_type="invalid",
@@ -437,8 +437,10 @@ def test_retrieve_documents(rag_retriever):
 # End-to-end functionality test
 def test_generate_answer(rag_retriever):
     """Test generating an answer using the RAG retriever."""
-    # Mock the retrieve method
-    rag_retriever.retrieve = Mock(
+    # Mock the retrieve method using patch
+    with patch.object(
+        rag_retriever,
+        "retrieve",
         return_value=(
             np.array([[[0.1] * 8, [0.2] * 8]]),  # 8-dimensional embeddings
             np.array([[1, 2]]),
@@ -449,28 +451,29 @@ def test_generate_answer(rag_retriever):
                     "title": ["Doc 1", "Doc 2"],
                 }
             ],
+        ),
+    ) as mock_retrieve:
+        # Mock the generator model's generate method
+        rag_retriever.generator_model.generate = Mock(
+            return_value=torch.tensor([[1, 2, 3]])
         )
-    )
-
-    # Mock the generator model's generate method
-    rag_retriever.generator_model.generate = Mock(
-        return_value=torch.tensor([[1, 2, 3]])
-    )
 
-    # Generate an answer
-    answer = rag_retriever.generate_answer("test query", top_k=2, max_new_tokens=100)
+        # Generate an answer
+        answer = rag_retriever.generate_answer(
+            "test query", top_k=2, max_new_tokens=100
+        )
 
-    # Verify the answer
-    assert isinstance(answer, str)
-    assert len(answer) > 0
+        # Verify the answer
+        assert isinstance(answer, str)
+        assert len(answer) > 0
 
-    # Verify that retrieve was called with correct parameters
-    rag_retriever.retrieve.assert_called_once()
-    call_args = rag_retriever.retrieve.call_args[1]
-    assert call_args["n_docs"] == 2
-    assert call_args["query"] == "test query"
+        # Verify that retrieve was called with correct parameters
+        mock_retrieve.assert_called_once()
+        call_args = mock_retrieve.call_args[1]
+        assert call_args["n_docs"] == 2
+        assert call_args["query"] == "test query"
 
-    # Verify that generate was called with correct parameters
-    rag_retriever.generator_model.generate.assert_called_once()
-    call_args = rag_retriever.generator_model.generate.call_args[1]
-    assert call_args["max_new_tokens"] == 100
+        # Verify that generate was called with correct parameters
+        rag_retriever.generator_model.generate.assert_called_once()
+        call_args = rag_retriever.generator_model.generate.call_args[1]
+        assert call_args["max_new_tokens"] == 100
diff --git a/sdk/python/tests/unit/test_vector_store.py b/sdk/python/tests/unit/test_vector_store.py
@@ -19,16 +19,16 @@ def cleanup_milvus():
     time.sleep(0.1)
 
 
-def test_vector_store_initialization(feature_store):
+def test_vector_store_initialization(example_feature_store):
     """Test vector store initialization."""
     print("Testing vector store initialization...")
 
     # Apply the feature view first
-    feature_store.apply([document_embeddings])
+    example_feature_store.apply([document_embeddings])
 
-    doc_view = feature_store.get_feature_view("document_embeddings")
+    doc_view = example_feature_store.get_feature_view("document_embeddings")
     store = MockVectorStore(
-        repo_path=str(feature_store.repo_path),
+        repo_path=str(example_feature_store.repo_path),
         rag_view=doc_view,
         features=[
             "document_embeddings:content",
@@ -44,16 +44,16 @@ def test_vector_store_initialization(feature_store):
     ]
 
 
-def test_vector_store_query(feature_store):
+def test_vector_store_query(example_feature_store):
     """Test vector store query method."""
     print("Testing vector store query...")
 
     # Apply the feature view first
-    feature_store.apply([document_embeddings])
+    example_feature_store.apply([document_embeddings])
 
-    doc_view = feature_store.get_feature_view("document_embeddings")
+    doc_view = example_feature_store.get_feature_view("document_embeddings")
     store = MockVectorStore(
-        repo_path=str(feature_store.repo_path),
+        repo_path=str(example_feature_store.repo_path),
         rag_view=doc_view,
         features=[
             "document_embeddings:content",
diff --git a/sdk/python/tests/utils/rag_test_utils.py b/sdk/python/tests/utils/rag_test_utils.py
@@ -71,7 +71,7 @@ def create_test_data():
 
 
 @pytest.fixture
-def feature_store():
+def example_feature_store():
     """Create a feature store using example repo."""
     runner = CliRunner()
     # Patch the run method to always succeed for teardown

Original file line number	Diff line number	Diff line change
`@@ -198,7 +198,6 @@ def test_dynamodb_online_store_update(repo_config, dynamodb_online_store):`
`198`	`198`	`assert existing_tables is not None`
`199`	`199`	`assert len(existing_tables) == 1`
`200`	`200`	`assert existing_tables[0] == f"test_aws.{db_table_keep_name}"`
`201`		`-`
`202`	`201`	`assert _get_tags(dynamodb_client, existing_tables[0]) == [`
`203`	`202`	`{"Key": "some", "Value": "tag"}`
`204`	`203`	`]`