test_modules.py

"""
🧪 تست کامل ماژول‌های ناشناس‌سازی
Complete testing suite for all modules
"""

import sys
import os

# اضافه کردن مسیر پروژه به Python path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from modules import (
    count_tokens,
    should_use_chunking,
    TextChunker,
    EntityNormalizer,
    EntityMerger
)


def test_utils():
    """تست توابع کمکی"""
    print("=" * 60)
    print("🧪 Test 1: Utils Module")
    print("=" * 60)
    
    # Token counting
    text = "این یک متن تست است برای شمارش توکن‌ها."
    tokens = count_tokens(text)
    print(f"\n✅ Token Counting:")
    print(f"  Text: {text}")
    print(f"  Tokens: {tokens}")
    assert tokens > 0, "Token count should be positive"
    
    # Chunking decision
    short_text = "متن کوتاه"
    long_text = "متن بلند " * 2000
    
    print(f"\n✅ Chunking Decision:")
    short_decision = should_use_chunking(short_text, threshold=1000)
    long_decision = should_use_chunking(long_text, threshold=1000)
    print(f"  Short text: {short_decision} (expected: False)")
    print(f"  Long text: {long_decision} (expected: True)")
    
    assert short_decision == False, "Short text should not need chunking"
    assert long_decision == True, "Long text should need chunking"
    
    print("\n✅ Utils tests passed!")


def test_normalizer():
    """تست نرمال‌سازی"""
    print("\n" + "=" * 60)
    print("🧪 Test 2: Normalizer Module")
    print("=" * 60)
    
    normalizer = EntityNormalizer()
    
    # Person
    print("\n✅ Person Normalization:")
    test_cases = [
        ("آقای علی احمدی", "علی احمدی"),
        ("دکتر مریم کریمی", "مریم کریمی"),
    ]
    
    for original, expected in test_cases:
        result = normalizer.normalize_person(original)
        print(f"  '{original}' → '{result}'")
        assert result == expected, f"Expected '{expected}', got '{result}'"
    
    # Company
    print("\n✅ Company Normalization:")
    company_result = normalizer.normalize_company("شرکت ملی نفت ایران")
    print(f"  'شرکت ملی نفت ایران' → '{company_result}'")
    assert "شرکت ملی نفت" in company_result
    
    # Amount
    print("\n✅ Amount Normalization:")
    amount_result = normalizer.normalize_amount("۵۰ میلیارد ریال")
    print(f"  '۵۰ میلیارد ریال' → '{amount_result}'")
    assert "50" in amount_result and "میلیارد" in amount_result
    
    # Percent
    print("\n✅ Percent Normalization:")
    percent_result = normalizer.normalize_percent("۱۵٪")
    print(f"  '۱۵٪' → '{percent_result}'")
    assert "15" in percent_result and "درصد" in percent_result
    
    print("\n✅ Normalizer tests passed!")


def test_chunker():
    """تست chunking"""
    print("\n" + "=" * 60)
    print("🧪 Test 3: Chunker Module")
    print("=" * 60)
    
    # متن تست
    test_text = """
    شرکت پارس در سال گذشته فروش 50 میلیارد ریال داشت. این رقم رشد 15 درصدی نشان می‌دهد.
    علی احمدی مدیرعامل شرکت است. شرکت صبا نیز در همین بازار فعالیت می‌کند.
    شرکت صبا فروش 30 میلیارد ریال داشت. مریم کریمی مدیر مالی است.
    همکاری بین دو شرکت در دستور کار است. قرارداد 20 میلیارد ریالی است.
    سرمایه‌گذاری 10 میلیارد ریال انجام می‌شود. بازده 25 درصد پیش‌بینی شده است.
    """
    
    chunker = TextChunker(chunk_size=100, overlap=20)
    chunks = chunker.create_chunks(test_text)
    
    print(f"\n✅ Chunking Results:")
    print(f"  Original text: {len(test_text)} chars")
    print(f"  Number of chunks: {len(chunks)}")
    
    for chunk in chunks:
        print(f"\n  {chunk['chunk_id']}:")
        print(f"    Tokens: {chunk['tokens']}")
        print(f"    Length: {chunk['length']} chars")
        print(f"    Preview: {chunk['text'][:60]}...")
    
    assert len(chunks) > 0, "Should create at least one chunk"
    assert chunker.validate_chunks(chunks), "Chunks should be valid"
    
    print("\n✅ Chunker tests passed!")


def test_merger():
    """تست merge"""
    print("\n" + "=" * 60)
    print("🧪 Test 4: Merger Module")
    print("=" * 60)
    
    merger = EntityMerger(fuzzy_threshold=0.75)
    
    # شبیه‌سازی دو mapping table
    table1 = {
        "chunk_id": "chunk_01",
        "mapping": {
            "company-01": "شرکت پارس",
            "company-02": "شرکت صبا",
            "person-01": "علی احمدی",
            "person-02": "مریم کریمی",
            "amount-01": "50 میلیارد ریال",
            "amount-02": "30 میلیارد ریال",
            "percent-01": "15 درصد"
        }
    }
    
    table2 = {
        "chunk_id": "chunk_02",
        "mapping": {
            "company-01": "شرکت  پارس",     # فاصله اضافی - باید match بشه
            "company-02": "شرکت ملی نفت",   # جدید
            "person-01": "علی احمدی",       # باید match بشه
            "person-02": "رضا محمدی",       # جدید
            "amount-01": "100 میلیارد ریال", # جدید
            "percent-01": "40 درصد"         # جدید
        }
    }
    
    result = merger.merge_mappings([table1, table2])
    
    print(f"\n✅ Merge Results:")
    print(f"\nGlobal Mapping ({len(result['global_mapping'])} entities):")
    for placeholder, value in sorted(result['global_mapping'].items()):
        print(f"  {placeholder}: {value}")
    
    print(f"\nRemapping for chunk_02:")
    remap = result['remapping'][1]['mapping']
    for old, new in remap.items():
        print(f"  {old} → {new}")
    
    # بررسی‌ها
    assert len(result['global_mapping']) > len(table1['mapping']), \
        "Global mapping should have more entities"
    
    assert len(result['remapping']) == 2, \
        "Should have remapping for 2 chunks"
    
    # بررسی exact match
    assert "company-01" in remap and remap["company-01"] == "company-01", \
        "شرکت پارس should match exactly"
    
    assert "person-01" in remap and remap["person-01"] == "person-01", \
        "علی احمدی should match exactly"
    
    print("\n✅ Merger tests passed!")


def test_end_to_end():
    """تست end-to-end کامل"""
    print("\n" + "=" * 60)
    print("🧪 Test 5: End-to-End Integration")
    print("=" * 60)
    
    # متن بلند برای تست
    long_text = """
    شرکت پارس در سال 1402 عملکرد بسیار خوبی داشت. فروش این شرکت به 50 میلیارد ریال رسید.
    علی احمدی، مدیرعامل شرکت پارس، اعلام کرد که این رشد 15 درصدی قابل توجه است.
    
    شرکت صبا نیز در همین صنعت فعالیت می‌کند. فروش شرکت صبا 30 میلیارد ریال بود.
    مریم کریمی، مدیر مالی شرکت صبا، پیش‌بینی کرد که در سال آینده به 40 میلیارد خواهد رسید.
    
    همکاری بین شرکت پارس و شرکت صبا در دستور کار قرار دارد. قرارداد به ارزش 20 میلیارد ریال
    در حال نهایی شدن است. علی احمدی و مریم کریمی در مذاکرات شرکت دارند.
    """
    
    # 1. تصمیم‌گیری
    use_chunking = should_use_chunking(long_text, threshold=50)
    print(f"\n📊 Use chunking: {use_chunking}")
    
    if use_chunking:
        # 2. Chunking
        chunker = TextChunker(chunk_size=100, overlap=20)
        chunks = chunker.create_chunks(long_text)
        print(f"✅ Created {len(chunks)} chunks")
        
        # 3. شبیه‌سازی ناشناس‌سازی (بدون Cerebras)
        # فرض: هر chunk یک mapping دارد
        mapping_tables = []
        
        for chunk in chunks:
            # شبیه‌سازی mapping
            mapping_tables.append({
                "chunk_id": chunk['chunk_id'],
                "mapping": {
                    "company-01": "شرکت پارس",
                    "person-01": "علی احمدی",
                    "amount-01": f"مبلغ تست {chunk['chunk_id']}"
                }
            })
        
        # 4. Merge
        merger = EntityMerger()
        merge_result = merger.merge_mappings(mapping_tables)
        
        print(f"✅ Merged to {len(merge_result['global_mapping'])} global entities")
        
        # 5. بررسی consistency
        # باید "شرکت پارس" و "علی احمدی" در همه chunks یکسان باشند
        for i, remap_data in enumerate(merge_result['remapping'][1:], start=2):
            remap = remap_data['mapping']
            if "company-01" in remap:
                assert remap["company-01"] == "company-01", \
                    f"company-01 in chunk {i} should map to global company-01"
        
        print("✅ Consistency check passed!")
    
    print("\n✅ End-to-end test passed!")


def run_all_tests():
    """اجرای تمام تست‌ها"""
    print("\n" + "🚀" * 30)
    print("Starting Complete Test Suite")
    print("🚀" * 30 + "\n")
    
    try:
        test_utils()
        test_normalizer()
        test_chunker()
        test_merger()
        test_end_to_end()
        
        print("\n" + "=" * 60)
        print("✅✅✅ ALL TESTS PASSED! ✅✅✅")
        print("=" * 60)
        
        return True
        
    except AssertionError as e:
        print("\n" + "=" * 60)
        print(f"❌ TEST FAILED: {e}")
        print("=" * 60)
        return False
    
    except Exception as e:
        print("\n" + "=" * 60)
        print(f"❌ ERROR: {e}")
        import traceback
        traceback.print_exc()
        print("=" * 60)
        return False


if __name__ == "__main__":
    success = run_all_tests()
    sys.exit(0 if success else 1)