test_integration.py

from db_manager import DBManager
from llm_client import LLMClient
import sys

def run_test():
    print("Running Integration Test...")
    
    # Setup
    try:
        db = DBManager()
        llm = LLMClient()
    except Exception as e:
        print(f"Setup Failed: {e}")
        sys.exit(1)

    # Test Case
    question = "How many products are in the Electronics category?"
    print(f"Test Question: {question}")
    
    try:
        # Schema
        print("Fetching Schema...")
        schema = db.get_schema_context()
        if not schema:
            print("Error: Empty schema returned.")
            sys.exit(1)
            
        # SQL Generation
        print("Generating SQL...")
        sql = llm.generate_sql(schema, question)
        print(f"Generated SQL: {sql}")
        
        if "SELECT" not in sql.upper():
            print("Error: Invalid SQL generated.")
            # Don't exit, might be LLM specific result, but let's log it
        
        # Execution
        print("Executing SQL...")
        results = db.execute_query(sql)
        
        print("Results:")
        print(results)
        
        if "data" in results:
            print("Test Passed: Query executed and returned data.")
        else:
            print("Test Failed: No data returned or error.")
            sys.exit(1)

    except Exception as e:
        print(f"Test Failed with Exception: {e}")
        sys.exit(1)
    finally:
        db.close()

if __name__ == "__main__":
    run_test()