68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
"""
|
||
Unit tests for the SearchAgent implementation.
|
||
"""
|
||
|
||
import unittest
|
||
|
||
from src.agent import SearchAgent, Result
|
||
|
||
|
||
class TestSearchAgent(unittest.TestCase):
|
||
def setUp(self):
|
||
self.corpus = [
|
||
"The quick brown fox jumps over the lazy dog.",
|
||
"Deep learning models can capture complex patterns in data.",
|
||
"PyTorch is a popular deep learning framework.",
|
||
"Natural language processing involves understanding text.",
|
||
]
|
||
self.agent = SearchAgent(self.corpus, embedding_dim=20)
|
||
|
||
def test_vocab_size(self):
|
||
# Vocabulary should contain all unique words
|
||
vocab_size = self.agent.get_vocab_size()
|
||
# Count unique words manually
|
||
unique_words = set()
|
||
for doc in self.corpus:
|
||
unique_words.update(doc.lower().split())
|
||
self.assertEqual(vocab_size, len(unique_words))
|
||
|
||
def test_document_embeddings_shape(self):
|
||
# Document embeddings should have shape (n_docs, dim)
|
||
doc_emb = self.agent.doc_embeddings
|
||
self.assertEqual(doc_emb.shape, (len(self.corpus), self.agent.embedding_dim))
|
||
|
||
def test_query_encoding_shape(self):
|
||
query = "deep learning"
|
||
vec = self.agent._encode_text(query)
|
||
self.assertEqual(vec.shape, (self.agent.embedding_dim,))
|
||
|
||
def test_cosine_similarity(self):
|
||
# Compute similarity between two identical vectors
|
||
vec = self.agent._encode_text("deep learning")
|
||
sims = self.agent._cosine_similarity(vec.unsqueeze(0), vec.unsqueeze(0))
|
||
self.assertAlmostEqual(sims.item(), 1.0, places=5)
|
||
|
||
def test_search_ranking(self):
|
||
# Query that matches second document
|
||
results = self.agent.search("deep learning", top_k=2)
|
||
# The first result should be the second document (index 1)
|
||
self.assertEqual(results[0].doc_id, 1)
|
||
self.assertGreater(results[0].score, results[1].score)
|
||
|
||
def test_unknown_words(self):
|
||
# Query with unknown words should still return results
|
||
results = self.agent.search("xyz abc", top_k=1)
|
||
self.assertEqual(len(results), 1)
|
||
self.assertIsInstance(results[0], Result)
|
||
|
||
def test_empty_query(self):
|
||
# Empty query should return top documents based on zero vector
|
||
results = self.agent.search("", top_k=3)
|
||
self.assertEqual(len(results), 3)
|
||
# Scores should be non‑negative
|
||
for res in results:
|
||
self.assertGreaterEqual(res.score, 0.0)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main() |