Spaces:
Sleeping
Sleeping
File size: 1,503 Bytes
e71fabd | 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 | import pytest
import tempfile
import os
from core.ingest import DocumentLoader, DocumentChunker, HierarchyManager
from core.utils import TextProcessor
class TestTextProcessor:
def test_count_tokens(self):
processor = TextProcessor()
text = "Hello world, this is a test."
tokens = processor.count_tokens(text)
assert tokens > 0
def test_mask_pii(self):
processor = TextProcessor()
text = "Contact me at test@example.com or 555-123-4567"
masked = processor.mask_pii(text)
assert "[EMAIL]" in masked
assert "[PHONE]" in masked
class TestDocumentLoader:
def test_load_txt(self):
loader = DocumentLoader()
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
f.write("Test content for text file.")
temp_path = f.name
try:
content = loader.load_txt(temp_path)
assert "Test content" in content
finally:
os.unlink(temp_path)
class TestHierarchyManager:
def test_load_hierarchies(self):
manager = HierarchyManager()
hierarchies = manager.list_hierarchies()
assert len(hierarchies) > 0
assert "hospital" in hierarchies
def test_get_hierarchy(self):
manager = HierarchyManager()
hospital_hierarchy = manager.get_hierarchy("hospital")
assert "levels" in hospital_hierarchy
assert "level1" in hospital_hierarchy["levels"] |