DronePlanning/tools/test_api.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import requests
import json
import os
from datetime import datetime

# --- Configuration ---
# The base URL of your running FastAPI service
BASE_URL = "http://127.0.0.1:8000"

# The endpoint we want to test
ENDPOINT = "/generate_plan"

# The user prompt we will send for the test
TEST_PROMPT = "已知目标检测红色气球危险性高于蓝色气球高于绿色气球，飞往搜索区搜索并锁定危险性最高的气球，对其跟踪30秒后进行打击操作"

# Log file path (will be created in the same directory as this script)
LOG_FILE = os.path.join(os.path.dirname(__file__), "api_test.log")

def write_log(message, print_to_console=True):
    """
    Write a message to the log file in append mode.
    Supports multi-line messages - only the first line gets timestamp.
    
    Args:
        message: The message to write (can be multi-line)
        print_to_console: Whether to also print to console (default: True)
    """
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # Split message into lines and add timestamp to first line only
    lines = message.split('\n')
    log_entries = [f"[{timestamp}] {lines[0]}\n"]
    for line in lines[1:]:
        log_entries.append(f"{' ' * (len(timestamp) + 3)}{line}\n")  # Indent continuation lines
    
    try:
        with open(LOG_FILE, "a", encoding="utf-8") as f:
            f.writelines(log_entries)
    except Exception as e:
        print(f"⚠️  Warning: Failed to write to log file: {e}")
    
    if print_to_console:
        print(message)

def test_generate_plan():
    """
    Sends a request to the /generate_plan endpoint and validates the response.
    All results are logged to the log file for continuous tracking.
    """
    url = BASE_URL + ENDPOINT
    payload = {"user_prompt": TEST_PROMPT}
    headers = {"Content-Type": "application/json"}

    # Write separator and test start info to log
    write_log("=" * 80, print_to_console=False)
    write_log("--- API Test: Generate Plan ---")
    write_log(f"URL: {url}")
    write_log(f"Sending Prompt: \"{TEST_PROMPT}\"")

    try:
        # Send the POST request
        response = requests.post(url, data=json.dumps(payload), headers=headers)
        
        # Check for HTTP errors (e.g., 404, 500)
        response.raise_for_status()

        # Parse the JSON response
        data = response.json()

        # Extract and log organized prompt if available in response
        organized_prompt = None
        if isinstance(data, dict):
            # Check for various possible field names for organized prompt
            organized_prompt = data.get("organized_prompt") or \
                             data.get("processed_prompt") or \
                             data.get("final_prompt") or \
                             data.get("enhanced_prompt") or \
                             data.get("user_prompt_enhanced")
        
        write_log("✅ Received Response:")
        
        # Log organized prompt if found
        if organized_prompt:
            write_log("\n📝 组织后的Prompt:")
            write_log(organized_prompt)
        else:
            # If not in response, log the original prompt for reference
            write_log("\n📝 原始Prompt:")
            write_log(f"   {TEST_PROMPT}")
            write_log("   (注: 组织后的prompt未在API响应中返回，如需查看请检查后端日志)")
        
        response_json = json.dumps(data, indent=2, ensure_ascii=False)
        write_log("\n完整响应内容:")
        write_log(response_json)

        # --- Validation ---
        write_log("\n--- Validation Checks ---")
        
        validation_results = []
        
        # 1. Check if the response is a dictionary
        if isinstance(data, dict):
            validation_results.append("PASS: Response is a valid JSON object.")
        else:
            validation_results.append("FAIL: Response is not a valid JSON object.")
            # Write all validation results to log before returning
            for result in validation_results:
                write_log(result)
            write_log("=" * 80, print_to_console=False)
            write_log("", print_to_console=False)  # Empty line for readability
            return

        # 2. Check for the existence of the 'root' key
        if "root" in data and isinstance(data['root'], dict):
            validation_results.append("PASS: Response contains a valid 'root' key.")
        else:
            validation_results.append("FAIL: Response does not contain a valid 'root' key.")

        # 3. Check for the existence and format of the 'visualization_url' key
        if "visualization_url" in data and data["visualization_url"].endswith(".png"):
            validation_results.append(f"PASS: Response contains a valid 'visualization_url': {data['visualization_url']}")
        else:
            validation_results.append("FAIL: Response does not contain a valid 'visualization_url'.")

        # Write all validation results to log
        for result in validation_results:
            write_log(result)

        # Write test completion marker
        write_log("✅ Test completed successfully")
        write_log("=" * 80, print_to_console=False)
        write_log("", print_to_console=False)  # Empty line for readability

    except requests.exceptions.RequestException as e:
        error_msg = f"❌ TEST FAILED: Could not connect to the server.\n   Please make sure the backend service is running.\n   Error details: {e}"
        write_log(error_msg)
        write_log("=" * 80, print_to_console=False)
        write_log("", print_to_console=False)  # Empty line for readability
    except json.JSONDecodeError:
        error_msg = f"❌ TEST FAILED: The server response was not valid JSON.\n   Response text: {response.text}"
        write_log(error_msg)
        write_log("=" * 80, print_to_console=False)
        write_log("", print_to_console=False)  # Empty line for readability
    except Exception as e:
        error_msg = f"❌ TEST FAILED: An unexpected error occurred: {e}"
        write_log(error_msg)
        write_log("=" * 80, print_to_console=False)
        write_log("", print_to_console=False)  # Empty line for readability

if __name__ == "__main__":
    test_generate_plan()