Autonomous-AI/test_runner.py at main · SimpNick6703/Autonomous-AI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import sys
import time
import httpx
import logging

logging.basicConfig(level=logging.INFO, format="[TEST] %(asctime)s - %(message)s")

BASE_URL = "http://localhost:8000"

TEST_CASES = [
    {
        "name": "Standard Request (PRD)",
        "request": "Create a comprehensive Product Requirements Document (PRD) for an AI-powered Customer Support Chatbot. Include target audience, key features, user flows, and technical integration details."
    },
    {
        "name": "Complex/Ambiguous Request (Incident SOP)",
        "request": "We had a major production incident last week with our payment gateway. Create the appropriate documentation for our engineering team."
    }
]


def run_tests():
    """Executes end-to-end async job submission and polling tests against the FastAPI server."""
    logging.info("Starting Autonomous Document Agent Test Suite...")

    with httpx.Client(timeout=30.0) as client:
        # Health Check
        try:
            resp = client.get(f"{BASE_URL}/docs")
            if resp.status_code != 200:
                logging.error("Server API docs not reachable. Is FastAPI running on port 8000?")
                sys.exit(1)
            logging.info("FastAPI server health check passed.")
        except Exception as e:
            logging.error(f"Could not connect to FastAPI server at {BASE_URL}: {e}")
            logging.info("Make sure to start the server first with: uvicorn main:app --port 8000")
            sys.exit(1)

        for tc in TEST_CASES:
            logging.info("=" * 70)
            logging.info(f"Running Test Case: {tc['name']}")
            logging.info(f"Request: '{tc['request']}'")
            logging.info("=" * 70)

            # POST /agent
            submit_resp = client.post(f"{BASE_URL}/agent", json={"request": tc["request"]})
            if submit_resp.status_code != 202:
                logging.error(f"Job submission failed: {submit_resp.status_code} - {submit_resp.text}")
                continue

            job_data = submit_resp.json()
            job_id = job_data["job_id"]
            logging.info(f"Job submitted successfully. Job ID: {job_id}, Status: {job_data['status']}")

            # Poll GET /agent/job/{job_id}
            poll_url = f"{BASE_URL}/agent/job/{job_id}"
            start_time = time.time()
            completed = False

            while time.time() - start_time < 300:  # 5 min timeout
                poll_resp = client.get(poll_url)
                if poll_resp.status_code != 200:
                    logging.error(f"Polling error: {poll_resp.status_code}")
                    break

                job_state = poll_resp.json()
                status = job_state["status"]
                current_step = job_state.get("current_step")

                logging.info(f"Polling Job {job_id} -> Status: {status} | Current Step: {current_step}")

                if status == "completed":
                    completed = True
                    logging.info(f"Job COMPLETED in {time.time() - start_time:.2f} seconds!")
                    logging.info(f"Generated Document Path: {job_state.get('document_path')}")
                    logging.info(f"Execution Plan Steps: {len(job_state.get('execution_plan', []))}")
                    break
                elif status == "failed":
                    logging.error(f"Job FAILED. Error: {job_state.get('error')}")
                    break

                time.sleep(3)

            if not completed:
                logging.error(f"Test case '{tc['name']}' did not complete within timeout.")


if __name__ == "__main__":
    run_tests()