test_agent_flow.py

"""Test script to verify the agent flow works correctly."""

from agent import agent
from database import db, Platform

def test_profile_collection():
    """Test the profile collection flow."""
    
    import time
    platform = "telegram"
    user_id = f"test_user_{int(time.time())}"  # Unique user ID
    
    print("=" * 60)
    print("Testing Profile Collection Flow")
    print(f"Test User ID: {user_id}")
    print("=" * 60)
    
    # Test 1: First query - should ask for location
    print("\n1. User says: Hi")
    result = agent.process_request(platform, user_id, "Hi")
    print(f"Bot: {result['response_text']}")
    assert "location" in result['response_text'].lower(), "Should ask for location"
    
    # Test 2: User provides location
    print("\n2. User says: Dindigul")
    result = agent.process_request(platform, user_id, "Dindigul")
    print(f"Bot: {result['response_text']}")
    assert "name" in result['response_text'].lower(), "Should ask for name"
    
    # Verify location is stored
    user = db.get_user(Platform(platform), user_id)
    print(f"   [OK] Location stored: {user.get('location_str')}")
    assert user.get('location_str') == "Dindigul", "Location should be stored"
    
    # Test 3: User provides name
    print("\n3. User says: Raj")
    result = agent.process_request(platform, user_id, "Raj")
    print(f"Bot: {result['response_text']}")
    assert "age" in result['response_text'].lower(), "Should ask for age"
    
    # Verify name is stored
    user = db.get_user(Platform(platform), user_id)
    print(f"   [OK] Name stored: {user.get('name')}")
    assert user.get('name') == "Raj", "Name should be stored"
    
    # Test 4: User provides age
    print("\n4. User says: 25")
    result = agent.process_request(platform, user_id, "25")
    print(f"Bot: {result['response_text'].encode('ascii', 'ignore').decode('ascii')}")
    assert "product" in result['response_text'].lower(), "Should mention product suggestions"
    
    # Verify age is stored
    user = db.get_user(Platform(platform), user_id)
    print(f"   [OK] Age stored: {user.get('age')}")
    assert user.get('age') == 25, "Age should be stored"
    
    # Test 5: Product query
    print("\n5. User says: I need a mobile under 5000")
    result = agent.process_request(platform, user_id, "I need a mobile under 5000")
    response_ascii = result['response_text'][:200].encode('ascii', 'ignore').decode('ascii')
    print(f"Bot: {response_ascii}...")
    print(f"   [OK] Picks returned: {len(result.get('picks', []))}")
    
    print("\n" + "=" * 60)
    print("[SUCCESS] All tests passed!")
    print("=" * 60)
    
    # Show final user profile
    user = db.get_user(Platform(platform), user_id)
    print("\nFinal User Profile:")
    print(f"  Name: {user.get('name')}")
    print(f"  Age: {user.get('age')}")
    print(f"  Location: {user.get('location_str')}")

if __name__ == "__main__":
    test_profile_collection()