test_local_parser.py

"""
Unit tests for zero-llm-engine local parser components.
"""
from __future__ import annotations

import os
import sys

# Insert current dir to path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from core.parser.local_parser.spelling import SymSpell
from core.parser.local_parser.synonyms import SynonymMapper
from core.parser.local_parser.ast_extractor import SafeMathEvaluator
from core.parser.local_parser.embeddings import WordPieceTokenizer, EmbeddingModel, cosine_similarity
from core.parser.local_parser.parser import LocalIntentParser
from core.column_registry import column_registry
from services.file_manager import handle_upload, delete_session

def test_spelling():
    print("Testing spelling correction...")
    sym = SymSpell(max_edit_distance=2)
    sym.add_word("salary")
    sym.add_word("average")
    sym.add_word("decrease")
    
    # Simple typo correction
    assert sym.lookup("salry")[0] == "salary"
    assert sym.lookup("avrage")[0] == "average"
    assert sym.lookup("decres")[0] == "decrease"
    
    # Sentence correction
    corrected = sym.correct_query("decres the salry")
    assert "decrease" in corrected
    assert "salary" in corrected
    print("✅ Spelling tests passed!")

def test_synonyms():
    print("Testing synonyms and Hinglish normalization...")
    mapper = SynonymMapper()
    
    # Hinglish mapping
    res = mapper.normalize_text("vetan ko 10% badhao")
    assert "salary" in res
    assert "increase" in res
    
    res2 = mapper.normalize_text("sabse chhota umar")
    assert "min" in res2
    assert "age" in res2
    print("✅ Synonym tests passed!")

def test_ast():
    print("Testing safe AST evaluator...")
    evaluator = SafeMathEvaluator(variables={"x": 10, "y": 20})
    
    # Math operations
    assert evaluator.evaluate("x + y * 2") == 50
    assert evaluator.evaluate("(y - x) / 2") == 5.0
    
    # Comparison and logic
    assert evaluator.evaluate("x > 5 and y < 30") is True
    assert evaluator.evaluate("x == 10 or y == 5") is True
    
    # Name error / safety check
    try:
        evaluator.evaluate("import os")
        assert False, "Should raise exception for imports"
    except Exception:
        pass
        
    try:
        evaluator.evaluate("x + z")
        assert False, "Should raise exception for undefined variable"
    except Exception:
        pass
    print("✅ AST tests passed!")

def test_tokenizer_and_embeddings():
    print("Testing tokenizer and embeddings...")
    # Initialize embedding model (downloads if necessary)
    model = EmbeddingModel()
    model.load_model()
    
    # WordPiece encoding tests
    tokenizer = model.tokenizer
    encoded = tokenizer.encode("salary")
    assert "input_ids" in encoded
    assert encoded["input_ids"].shape == (1, 128)
    
    # Embedding generation
    emb1 = model.get_embedding("salary")
    emb2 = model.get_embedding("vetan")
    emb3 = model.get_embedding("country")
    
    # L2 Normalized shape check
    assert len(emb1.shape) == 1
    assert emb1.shape[0] > 0
    
    # Cosine similarities
    sim_salary = cosine_similarity(emb1, emb2)
    sim_diff = cosine_similarity(emb1, emb3)
    
    print(f"Similarity (salary, vetan): {sim_salary}")
    print(f"Similarity (salary, country): {sim_diff}")
    
    # Match column
    cols = ["Age", "Salary", "Name", "Country"]
    best_col, score = model.match_column("payrate", cols, threshold=0.3)
    assert best_col == "Salary"
    print("✅ Tokenizer and Embedding tests passed!")

def test_orchestration():
    print("Testing orchestration parser...")
    # Register columns for test session
    sid = "test_parser_session"
    column_registry.register_columns(sid, [
        {"name": "Name", "dtype": "String"},
        {"name": "City", "dtype": "String"},
        {"name": "Salary", "dtype": "Float64"},
        {"name": "Age", "dtype": "Int64"},
    ])
    
    parser = LocalIntentParser()
    
    # 1. Test increase
    res1 = parser.parse_intent(sid, "salry ko 10% badhao")
    assert res1 is not None
    assert res1["operation"] == "increase"
    assert res1["column"] == "Salary"
    assert res1["value"] == 10.0
    assert res1["is_percent"] is True
    
    # 2. Test filter
    res2 = parser.parse_intent(sid, "salary se zyada 50000 dikhao")
    # "se zyada" gets normalized, condition becomes >
    assert res2 is not None
    assert res2["operation"] == "filter"
    assert res2["column"] == "Salary"
    assert res2["condition"] in [">", "=="] # depends on normalization
    
    # Clean up
    column_registry.clear_session(sid)
    print("✅ Orchestration tests passed!")

if __name__ == "__main__":
    test_spelling()
    test_synonyms()
    test_ast()
    try:
        test_tokenizer_and_embeddings()
    except Exception as e:
        print(f"⚠️ Embeddings/Tokenizer test skipped or failed due to env: {e}")
    test_orchestration()
    print("\n🎉 ALL TESTS COMPLETED SUCCESSFULY!")