File size: 10,797 Bytes
f28b6fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 |
"""
Tests for FinEE RAG Engine
==========================
Comprehensive tests for:
- Merchant Knowledge Base
- Category Taxonomy
- Vector Store
- RAG Engine
Author: Ranjit Behera
"""
import pytest
import json
from pathlib import Path
import tempfile
class TestMerchantKnowledgeBase:
"""Tests for MerchantKnowledgeBase."""
def setup_method(self):
"""Setup test fixtures."""
from finee.rag import MerchantKnowledgeBase, Merchant
self.kb = MerchantKnowledgeBase()
self.Merchant = Merchant
def test_default_merchants_loaded(self):
"""Test that default merchants are loaded."""
assert len(self.kb.merchants) > 30
assert "swiggy" in self.kb.merchants
assert "amazon" in self.kb.merchants
assert "zerodha" in self.kb.merchants
def test_lookup_by_name(self):
"""Test merchant lookup by name."""
merchant = self.kb.lookup("Swiggy")
assert merchant is not None
assert merchant.name == "Swiggy"
assert merchant.category == "food"
def test_lookup_by_vpa(self):
"""Test merchant lookup by VPA."""
merchant = self.kb.lookup("swiggy@ybl")
assert merchant is not None
assert merchant.name == "Swiggy"
def test_lookup_by_alias(self):
"""Test merchant lookup by alias."""
merchant = self.kb.lookup("amzn")
assert merchant is not None
assert merchant.name == "Amazon"
def test_lookup_partial_match(self):
"""Test partial name matching."""
merchant = self.kb.lookup("netflix")
assert merchant is not None
assert merchant.name == "Netflix"
def test_lookup_not_found(self):
"""Test lookup returns None for unknown merchant."""
merchant = self.kb.lookup("unknownmerchant123")
assert merchant is None
def test_search_by_text(self):
"""Test text-based merchant search."""
matches = self.kb.search("food delivery swiggy order")
assert len(matches) > 0
assert matches[0].name == "Swiggy"
def test_add_merchant(self):
"""Test adding new merchant."""
new_merchant = self.Merchant(
name="TestMerchant",
vpa="test@ybl",
category="test",
aliases=["tm"],
)
self.kb.add_merchant(new_merchant)
assert "testmerchant" in self.kb.merchants
assert self.kb.lookup("test@ybl") is not None
assert self.kb.lookup("tm") is not None
def test_get_category_merchants(self):
"""Test getting merchants by category."""
food_merchants = self.kb.get_category_merchants("food")
assert len(food_merchants) > 0
for m in food_merchants:
assert m.category == "food"
def test_save_and_load(self):
"""Test saving and loading KB."""
with tempfile.TemporaryDirectory() as tmpdir:
path = Path(tmpdir) / "merchants.json"
# Save
self.kb.save(path)
assert path.exists()
# Load
from finee.rag import MerchantKnowledgeBase
loaded_kb = MerchantKnowledgeBase.load(path)
assert len(loaded_kb.merchants) == len(self.kb.merchants)
class TestCategoryTaxonomy:
"""Tests for CategoryTaxonomy."""
def setup_method(self):
"""Setup test fixtures."""
from finee.rag import CategoryTaxonomy
self.taxonomy = CategoryTaxonomy
def test_get_hierarchy_simple(self):
"""Test simple category hierarchy."""
hierarchy = self.taxonomy.get_hierarchy("food")
assert hierarchy == ["food"]
def test_get_hierarchy_nested(self):
"""Test nested category hierarchy."""
hierarchy = self.taxonomy.get_hierarchy("grocery")
assert "shopping" in hierarchy
assert "grocery" in hierarchy
def test_get_hierarchy_unknown(self):
"""Test unknown category returns itself."""
hierarchy = self.taxonomy.get_hierarchy("unknown")
assert hierarchy == ["unknown"]
def test_infer_category_food(self):
"""Test food category inference."""
category = self.taxonomy.infer_category("lunch delivery from restaurant")
assert category == "food"
def test_infer_category_investment(self):
"""Test investment category inference."""
category = self.taxonomy.infer_category("SIP mutual fund trading invest")
assert category == "investment"
def test_infer_category_bills(self):
"""Test bills category inference."""
category = self.taxonomy.infer_category("electricity bill payment recharge")
assert category == "bills"
def test_infer_category_unknown(self):
"""Test unknown text returns 'other'."""
category = self.taxonomy.infer_category("random text here")
assert category == "other"
class TestSimpleVectorStore:
"""Tests for SimpleVectorStore."""
def setup_method(self):
"""Setup test fixtures."""
from finee.rag import SimpleVectorStore, Transaction
self.store = SimpleVectorStore()
self.Transaction = Transaction
def test_add_transaction(self):
"""Test adding transaction."""
txn = self.Transaction(
id="test1",
text="HDFC Bank Rs.500 debited Swiggy",
amount=500.0,
type="debit",
merchant="Swiggy",
category="food",
)
self.store.add(txn)
assert len(self.store.documents) == 1
def test_search_similar(self):
"""Test similarity search."""
# Add some transactions
transactions = [
self.Transaction("1", "HDFC Rs.500 Swiggy food", 500, "debit", "Swiggy", "food"),
self.Transaction("2", "SBI Rs.1000 Amazon shopping", 1000, "debit", "Amazon", "shopping"),
self.Transaction("3", "ICICI Rs.250 Zomato dinner", 250, "debit", "Zomato", "food"),
]
for txn in transactions:
self.store.add(txn)
# Search
results = self.store.search("Swiggy food delivery", limit=2)
assert len(results) > 0
# Swiggy or Zomato should be top result (both are food)
assert results[0][0].category == "food"
def test_search_empty_store(self):
"""Test search on empty store."""
results = self.store.search("any query")
assert results == []
def test_save_and_load(self):
"""Test saving and loading store."""
# Add transactions
txn = self.Transaction("1", "Test transaction", 100, "debit")
self.store.add(txn)
with tempfile.TemporaryDirectory() as tmpdir:
path = Path(tmpdir) / "store.json"
# Save
self.store.save(path)
assert path.exists()
# Load
from finee.rag import SimpleVectorStore
loaded = SimpleVectorStore.load(path)
assert len(loaded.documents) == 1
class TestRAGEngine:
"""Tests for RAGEngine."""
def setup_method(self):
"""Setup test fixtures."""
from finee.rag import RAGEngine
self.rag = RAGEngine()
def test_retrieve_with_merchant(self):
"""Test retrieval with known merchant."""
context = self.rag.retrieve("HDFC Bank Rs.499 debited UPI:swiggy@ybl")
assert context.merchant_info is not None
assert context.merchant_info["name"] == "Swiggy"
assert context.merchant_info["category"] == "food"
assert context.confidence_boost > 0
def test_retrieve_with_category(self):
"""Test category inference."""
context = self.rag.retrieve("Netflix subscription payment")
assert context.category_hierarchy is not None
assert "entertainment" in context.category_hierarchy
def test_retrieve_investment(self):
"""Test investment detection."""
context = self.rag.retrieve("Rs.25000 transferred to Zerodha trading")
assert context.merchant_info is not None
assert context.merchant_info["category"] == "investment"
def test_augment_prompt(self):
"""Test prompt augmentation."""
context = self.rag.retrieve("Swiggy food order")
augmented = self.rag.augment_prompt("Swiggy food order", context)
assert "Swiggy" in augmented
assert "food" in augmented
def test_add_transaction(self):
"""Test adding transaction to history."""
self.rag.add_transaction(
"Test transaction text",
{"amount": 100, "merchant": "Test", "category": "test"}
)
assert len(self.rag.vector_store.documents) > 0
def test_save_and_load(self):
"""Test saving and loading RAG state."""
with tempfile.TemporaryDirectory() as tmpdir:
path = Path(tmpdir) / "rag_state"
# Save
self.rag.save(path)
assert (path / "merchants.json").exists()
# Load
from finee.rag import RAGEngine
new_rag = RAGEngine()
new_rag.load(path)
assert len(new_rag.merchant_kb.merchants) > 0
class TestRAGIntegration:
"""Integration tests for RAG with extraction."""
def setup_method(self):
"""Setup test fixtures."""
from finee.rag import RAGEngine
self.rag = RAGEngine()
@pytest.mark.parametrize("message,expected_merchant,expected_category", [
("HDFC Rs.499 UPI:swiggy@ybl", "Swiggy", "food"),
("SBI Rs.999 Netflix subscription", "Netflix", "entertainment"),
("ICICI Rs.25000 Zerodha trading", "Zerodha", "investment"),
("Axis Rs.1500 Amazon order", "Amazon", "shopping"),
("Kotak Rs.350 Uber ride", "Uber", "transport"),
("PNB Rs.100 Airtel recharge", "Airtel", "bills"),
])
def test_merchant_detection(self, message, expected_merchant, expected_category):
"""Test merchant detection for various messages."""
context = self.rag.retrieve(message)
assert context.merchant_info is not None
assert context.merchant_info["name"] == expected_merchant
assert context.merchant_info["category"] == expected_category
# ============================================================================
# RUN TESTS
# ============================================================================
if __name__ == "__main__":
pytest.main([__file__, "-v"])
|