learningcircuit--local-deep-research/tests/ui_tests/test_simple_research_api.py

#!/usr/bin/env python3
"""Test research submission directly via API.

This is a manual utility script, not a pytest test.
Run directly with: python test_simple_research_api.py
"""

import json
import sys
import time

import requests


def main():
    """Test research submission via API."""
    # Login first
    session = requests.Session()
    login_data = {
        "username": "testuser_api",
        "password": "T3st!Secure#2024$LDR",
    }

    # Register if needed
    reg_page = session.get("http://127.0.0.1:5000/auth/register")
    csrf_token = reg_page.text.split('name="csrf_token" value="')[1].split('"')[
        0
    ]

    reg_data = login_data.copy()
    reg_data["confirm_password"] = login_data["password"]
    reg_data["csrf_token"] = csrf_token
    reg_data["acknowledge"] = "on"

    reg_resp = session.post(
        "http://127.0.0.1:5000/auth/register", data=reg_data
    )
    print(f"Registration status: {reg_resp.status_code}")

    # Login
    login_page = session.get("http://127.0.0.1:5000/auth/login")
    csrf_token = login_page.text.split('name="csrf_token" value="')[1].split(
        '"'
    )[0]
    login_data["csrf_token"] = csrf_token
    login_resp = session.post(
        "http://127.0.0.1:5000/auth/login", data=login_data
    )
    print(f"Login status: {login_resp.status_code}")

    # Get CSRF token for API
    home_page = session.get("http://127.0.0.1:5000/")
    if 'name="csrf-token"' in home_page.text:
        csrf_token = home_page.text.split('name="csrf-token" content="')[
            1
        ].split('"')[0]
    else:
        print("Not logged in properly, trying again...")
        # Try login with existing user
        login_data["username"] = "testuser"
        login_page = session.get("http://127.0.0.1:5000/auth/login")
        csrf_token = login_page.text.split('name="csrf_token" value="')[
            1
        ].split('"')[0]
        login_data["csrf_token"] = csrf_token
        login_resp = session.post(
            "http://127.0.0.1:5000/auth/login", data=login_data
        )
        print(f"Re-login status: {login_resp.status_code}")

        home_page = session.get("http://127.0.0.1:5000/")
        if 'name="csrf_token"' in home_page.text:
            csrf_token = home_page.text.split('name="csrf_token" value="')[
                1
            ].split('"')[0]
        else:
            print("Failed to get CSRF token")
            sys.exit(1)

    # Submit research
    research_data = {
        "query": "What is Python programming?",
        "mode": "quick",
        "model_provider": "OLLAMA",
        "model": "llama2",
        "search_engine": "searxng",
        "iterations": 1,
        "questions_per_iteration": 2,
        "strategy": "source-based",
    }

    headers = {"Content-Type": "application/json", "X-CSRFToken": csrf_token}

    print("\nSubmitting research...")
    submit_resp = session.post(
        "http://127.0.0.1:5000/api/start_research",
        json=research_data,
        headers=headers,
    )
    print(f"Submit status: {submit_resp.status_code}")

    if submit_resp.ok:
        result = submit_resp.json()
        print(f"Result: {json.dumps(result, indent=2)}")

        if result.get("status") == "success":
            research_id = result.get("research_id")
            print(f"\nResearch ID: {research_id}")

            # Check status after a few seconds
            for i in range(5):
                time.sleep(2)  # allow: unmarked-sleep
                status_resp = session.get(
                    f"http://127.0.0.1:5000/api/research/{research_id}/status"
                )
                if status_resp.ok:
                    status = status_resp.json()
                    print(
                        f"\nCheck {i + 1}: Status={status.get('status')}, "
                        f"Progress={status.get('progress')}"
                    )

                    # Get logs
                    logs_resp = session.get(
                        f"http://127.0.0.1:5000/api/research/{research_id}/logs"
                    )
                    if logs_resp.ok:
                        logs = logs_resp.json()
                        if isinstance(logs, list) and logs:
                            print(
                                f"Latest log: {logs[-1].get('message', 'No message')}"
                            )
    else:
        print(f"Error: {submit_resp.text[:500]}")


if __name__ == "__main__":
    main()