test_api.py

#!/usr/bin/env python
"""
Quick demo script to test the API endpoints
Run: python test_api.py
"""

import requests
import json
from typing import Dict, Any

BASE_URL = "http://localhost:8000"

# Colors for terminal output
GREEN = '\033[92m'
BLUE = '\033[94m'
YELLOW = '\033[93m'
RED = '\033[91m'
RESET = '\033[0m'


def print_header(title: str):
    """Print a formatted header"""
    print(f"\n{BLUE}{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}{RESET}\n")


def print_success(msg: str):
    """Print success message"""
    print(f"{GREEN}✓ {msg}{RESET}")


def print_error(msg: str):
    """Print error message"""
    print(f"{RED}✗ {msg}{RESET}")


def print_response(response: Dict[Any, Any]):
    """Pretty print JSON response"""
    print(json.dumps(response, indent=2))


def test_health_check():
    """Test health endpoint"""
    print_header("1. Health Check")
    try:
        response = requests.get(f"{BASE_URL}/health")
        print_success("Health check passed")
        print_response(response.json())
        return True
    except Exception as e:
        print_error(f"Health check failed: {e}")
        return False


def test_create_task():
    """Test task creation"""
    print_header("2. Create a Task")
    try:
        payload = {
            "title": "Review API documentation",
            "description": "Check the new endpoints for Slack integration"
        }
        response = requests.post(f"{BASE_URL}/tasks/", json=payload)
        if response.status_code == 200:
            print_success("Task created")
            print_response(response.json())
            return True
        else:
            print_error(f"Failed with status {response.status_code}")
            print_response(response.json())
            return False
    except Exception as e:
        print_error(f"Task creation failed: {e}")
        return False


def test_list_tasks():
    """Test listing tasks"""
    print_header("3. List All Tasks")
    try:
        response = requests.get(f"{BASE_URL}/tasks/")
        if response.status_code == 200:
            print_success("Tasks retrieved")
            print_response(response.json())
            return True
        else:
            print_error(f"Failed with status {response.status_code}")
            return False
    except Exception as e:
        print_error(f"Listing tasks failed: {e}")
        return False


def test_summarize():
    """Test text summarization"""
    print_header("4. Summarize Text")
    try:
        payload = {
            "text": """
            In today's meeting we discussed the Q1 roadmap. 
            John mentioned we need to prioritize the API refactoring. 
            Sarah will handle the frontend redesign. 
            The deployment is planned for end of month. 
            We also need to schedule a code review session.
            """,
            "store_tasks": True
        }
        response = requests.post(f"{BASE_URL}/summarize", json=payload)
        if response.status_code == 200:
            print_success("Text summarized")
            print_response(response.json())
            return True
        else:
            print_error(f"Failed with status {response.status_code}")
            print_response(response.json())
            return False
    except Exception as e:
        print_error(f"Summarization failed: {e}")
        return False


def test_ask_question():
    """Test asking a question"""
    print_header("5. Ask a Question (with context)")
    try:
        payload = {
            "user_id": "U123456789",
            "question": "What's on my plate today? What should I prioritize?",
            "thread_ts": "1234567.890123"
        }
        response = requests.post(f"{BASE_URL}/chat/ask", json=payload)
        if response.status_code == 200:
            print_success("Question answered")
            print_response(response.json())
            return True
        else:
            print_error(f"Failed with status {response.status_code}")
            print_response(response.json())
            return False
    except Exception as e:
        print_error(f"Question answering failed: {e}")
        return False


def test_pending_summary():
    """Test getting pending summary"""
    print_header("6. Get Pending Summary for User")
    try:
        response = requests.get(
            f"{BASE_URL}/chat/pending-summary/U123456789"
        )
        if response.status_code == 200:
            print_success("Pending summary retrieved")
            print_response(response.json())
            return True
        else:
            print_error(f"Failed with status {response.status_code}")
            print_response(response.json())
            return False
    except Exception as e:
        print_error(f"Pending summary failed: {e}")
        return False


def test_mark_thread():
    """Test marking a thread"""
    print_header("7. Mark Slack Thread for Analysis")
    try:
        payload = {
            "thread_ts": "1234567.890123",
            "is_marked": True
        }
        response = requests.post(
            f"{BASE_URL}/memory/slack-threads/mark",
            json=payload
        )
        if response.status_code == 200:
            print_success("Thread marked")
            print_response(response.json())
            return True
        else:
            print_error(f"Failed with status {response.status_code}")
            print_response(response.json())
            return False
    except Exception as e:
        print_error(f"Marking thread failed: {e}")
        return False


def test_get_marked_threads():
    """Test getting marked threads"""
    print_header("8. Get All Marked Threads")
    try:
        response = requests.get(
            f"{BASE_URL}/memory/slack-threads/marked"
        )
        if response.status_code == 200:
            print_success("Marked threads retrieved")
            print_response(response.json())
            return True
        else:
            print_error(f"Failed with status {response.status_code}")
            return False
    except Exception as e:
        print_error(f"Getting marked threads failed: {e}")
        return False


def test_get_user_threads():
    """Test getting user threads"""
    print_header("9. Get User Slack Threads")
    try:
        response = requests.get(
            f"{BASE_URL}/memory/slack-threads/user/U123456789"
        )
        if response.status_code == 200:
            print_success("User threads retrieved")
            print_response(response.json())
            return True
        else:
            print_error(f"Failed with status {response.status_code}")
            return False
    except Exception as e:
        print_error(f"Getting user threads failed: {e}")
        return False


def main():
    """Run all tests"""
    print(f"\n{YELLOW}{'='*60}")
    print("  AI Assistant API - Test Suite")
    print(f"{'='*60}{RESET}")
    print(f"\nTesting API at: {BASE_URL}")
    print("Make sure the server is running: uvicorn api.main:app --reload\n")

    results = []

    # Run tests
    results.append(("Health Check", test_health_check()))
    results.append(("Create Task", test_create_task()))
    results.append(("List Tasks", test_list_tasks()))
    results.append(("Summarize Text", test_summarize()))
    results.append(("Ask Question", test_ask_question()))
    results.append(("Pending Summary", test_pending_summary()))
    results.append(("Mark Thread", test_mark_thread()))
    results.append(("Get Marked Threads", test_get_marked_threads()))
    results.append(("Get User Threads", test_get_user_threads()))

    # Print summary
    print_header("Test Summary")
    passed = sum(1 for _, result in results if result)
    total = len(results)

    for test_name, result in results:
        status = f"{GREEN}PASSED{RESET}" if result else f"{RED}FAILED{RESET}"
        print(f"  {test_name}: {status}")

    print(f"\n{BLUE}Total: {passed}/{total} tests passed{RESET}\n")

    if passed == total:
        print_success("All tests passed! ✨")
    else:
        print_error(f"{total - passed} test(s) failed")


if __name__ == "__main__":
    main()