tests.py

# tests.py - Pruebas para verificar el funcionamiento

from client import LlamaClient
import time
import json

def test_basic_functionality(base_url: str):
    """Pruebas básicas de funcionalidad"""
    print("🧪 Ejecutando pruebas básicas...")
    client = LlamaClient(base_url)

    tests = []

    # Test 1: Chat simple
    print("\n1. Test chat simple...")
    try:
        response = client.chat(
            message="Hola, ¿puedes presentarte en una línea?",
            max_tokens=50
        )

        if "error" in response:
            tests.append(("Chat simple", False, response['error']))
        else:
            tests.append(("Chat simple", True, f"Respuesta: {response['response'][:50]}..."))
    except Exception as e:
        tests.append(("Chat simple", False, str(e)))

    # Test 2: Chat con system prompt
    print("2. Test system prompt...")
    try:
        response = client.chat(
            message="¿Cuánto es 2+2?",
            system_prompt="Eres una calculadora. Solo responde con números.",
            max_tokens=20
        )

        success = "error" not in response and "4" in response.get('response', '')
        tests.append(("System prompt", success, response.get('response', 'No response')))
    except Exception as e:
        tests.append(("System prompt", False, str(e)))

    # Test 3: Chat con historial
    print("3. Test historial...")
    try:
        history = [["¿Cómo te llamas?", "Soy un asistente de IA."]]
        response = client.chat(
            message="¿Recuerdas cómo te llamas?",
            history=history,
            max_tokens=50
        )

        success = "error" not in response
        tests.append(("Historial", success, response.get('response', 'Error')[:50]))
    except Exception as e:
        tests.append(("Historial", False, str(e)))

    # Test 4: Streaming
    print("4. Test streaming...")
    try:
        chunks_received = 0
        final_response = ""

        for chunk in client.chat_stream(
            message="Cuenta del 1 al 5",
            max_tokens=30
        ):
            chunks_received += 1
            if "error" in chunk:
                tests.append(("Streaming", False, chunk['error']))
                break

            final_response = chunk['response']
            if chunk.get("is_complete", False):
                break

        success = chunks_received > 1 and final_response
        tests.append(("Streaming", success, f"{chunks_received} chunks, respuesta: {final_response[:30]}"))
    except Exception as e:
        tests.append(("Streaming", False, str(e)))

    # Test 5: Estado de cola
    print("5. Test estado de cola...")
    try:
        status = client.get_queue_status()
        success = "queue_size" in status and "is_processing" in status
        tests.append(("Estado cola", success, str(status)))
    except Exception as e:
        tests.append(("Estado cola", False, str(e)))

    # Mostrar resultados
    print("\n" + "="*60)
    print("📊 RESULTADOS DE PRUEBAS")
    print("="*60)

    passed = 0
    for test_name, success, details in tests:
        status = "✅ PASS" if success else "❌ FAIL"
        print(f"{status} {test_name}: {details}")
        if success:
            passed += 1

    print(f"\n🎯 Resultado: {passed}/{len(tests)} pruebas pasaron")
    return passed == len(tests)

def stress_test(base_url: str, num_requests: int = 5):
    """Prueba de estrés con múltiples requests"""
    print(f"\n⚡ Prueba de estrés ({num_requests} requests)...")
    client = LlamaClient(base_url)

    start_time = time.time()
    results = []

    for i in range(num_requests):
        print(f"Request {i+1}/{num_requests}...", end=" ")

        try:
            request_start = time.time()
            response = client.chat(
                message=f"Esta es la request número {i+1}. Responde brevemente.",
                max_tokens=50
            )
            request_time = time.time() - request_start

            if "error" in response:
                print(f"❌ Error: {response['error']}")
                results.append({"success": False, "time": request_time, "error": response['error']})
            else:
                print(f"✅ {request_time:.2f}s")
                results.append({"success": True, "time": request_time, "response_length": len(response['response'])})

        except Exception as e:
            print(f"❌ Exception: {e}")
            results.append({"success": False, "time": 0, "error": str(e)})

    total_time = time.time() - start_time

    # Análisis de resultados
    successful = [r for r in results if r['success']]
    failed = [r for r in results if not r['success']]

    print(f"\n📈 Análisis de estrés:")
    print(f"  • Total: {total_time:.2f}s")
    print(f"  • Exitosas: {len(successful)}/{num_requests}")
    print(f"  • Fallidas: {len(failed)}/{num_requests}")

    if successful:
        avg_time = sum(r['time'] for r in successful) / len(successful)
        print(f"  • Tiempo promedio: {avg_time:.2f}s")

    return len(successful) == num_requests

def performance_benchmark(base_url: str):
    """Benchmark de rendimiento"""
    print("\n🏁 Benchmark de rendimiento...")
    client = LlamaClient(base_url)

    test_cases = [
        ("Respuesta corta", "Hola", 20),
        ("Respuesta media", "Explica qué es Python en un párrafo", 100),
        ("Respuesta larga", "Describe la historia de la programación", 300),
    ]

    for test_name, message, max_tokens in test_cases:
        print(f"\n{test_name} ({max_tokens} tokens)...")

        # Test sin streaming
        start_time = time.time()
        response = client.chat(message=message, max_tokens=max_tokens)
        normal_time = time.time() - start_time

        if "error" in response:
            print(f"  ❌ Error: {response['error']}")
            continue

        # Test con streaming
        start_time = time.time()
        for chunk in client.chat_stream(message=message, max_tokens=max_tokens):
            if chunk.get("is_complete", False):
                break
        stream_time = time.time() - start_time

        print(f"  • Sin streaming: {normal_time:.2f}s")
        print(f"  • Con streaming: {stream_time:.2f}s")
        print(f"  • Tokens generados: ~{len(response['response'])} chars")

def main():
    """Ejecutar todas las pruebas"""
    base_url = input("Ingresa la URL de tu Space: ").strip()

    if not base_url:
        print("❌ URL requerida")
        return

    if not base_url.startswith(('http://', 'https://')):
        base_url = f"https://{base_url}"

    print(f"🧪 Probando Space: {base_url}")
    print("="*60)

    # Ejecutar pruebas
    basic_ok = test_basic_functionality(base_url)

    if basic_ok:
        print("\n✅ Pruebas básicas exitosas. Continuando...")
        stress_ok = stress_test(base_url, 3)
        performance_benchmark(base_url)

        if stress_ok:
            print("\n🎉 ¡Todas las pruebas pasaron! El Space funciona correctamente.")
        else:
            print("\n⚠️ Algunas pruebas de estrés fallaron. Verifica el rendimiento.")
    else:
        print("\n❌ Pruebas básicas fallaron. Verifica la configuración del Space.")

if __name__ == "__main__":
    main()