test.py

import requests
import json
import re
from fastapi import HTTPException

from app import process_chunks, analysis_system_message, compute_input_size

url = "https://odai0-silmaq5.hf.space"


# /prompt
def test_prompt_endpoint():
    data = {"prompt": "say hi back"}
    response = requests.post(f"{url}/prompt", json=data)

    if response.status_code == 200:
        try:
            result = response.json()
            print("Valid JSON Response:", json.dumps(result, indent=2))
        except json.JSONDecodeError:
            print("Not valid JSON:", response.text)
            raise AssertionError("Invalid JSON response")

        output = result["output"]
        assert isinstance(output, dict), "Output is not a dictionary"
        assert "signal" in output, "Missing 'signal' key"
        assert "message" in output, "Missing 'message' key"
        assert output["signal"] in ["m0", "a0", "e0"], "Invalid 'signal' value"
        assert re.search(r"\bhi\b", output["message"], re.IGNORECASE), "Model didn't say hi"

    else:
        print(f"Request failed with status code: {response.status_code}")
        raise HTTPException(status_code=response.status_code, detail="Request to /prompt failed")

    # check for analysis requests
    dataArr = [
        {"prompt": "analyze this site"},
        # {"prompt": "what options are there"},
        {"prompt": "summarize the content"},
    ]

    for data in dataArr:
        response = requests.post(f"{url}/prompt", json=data)

        if response.status_code == 200:
            try:
                result = response.json()
                print("Valid JSON Response:", json.dumps(result, indent=2))
            except json.JSONDecodeError:
                print("Not valid JSON:", response.text)
                raise AssertionError("Invalid JSON response")

            output = result["output"]
            assert isinstance(output, dict), "Output is not a dictionary"
            assert "signal" in output, "Missing 'signal' key"
            assert "message" in output, "Missing 'message' key"
            assert output["signal"] == "a1", "Invalid 'signal' value, expected 'a1'"

        else:
            print(f"Request failed with status code: {response.status_code}")
            raise HTTPException(status_code=response.status_code, detail="Request to /prompt failed")

    print("Prompt endpoint test passed.")


# /analyze
def test_analyze_endpoint():
    chunk_str = json.dumps({
        "tag": "nav",
        "id": "",
        "elements": [
            {
                "tag": "ul",
                "id": "",
                "elements": [
                    {"tag": "li", "id": "", "elements": [{"tag": "a", "id": "", "elements": [{"tag": "i", "id": "", "elements": []}, {"id": "", "tag": "#text", "text": "Category"}]}]},
                    {"tag": "li", "id": "", "elements": [{"tag": "a", "id": "", "elements": [{"tag": "i", "id": "", "elements": []}, {"id": "", "tag": "#text", "text": "Articles"}]}]},
                    {"tag": "li", "id": "", "elements": [{"tag": "a", "id": "", "elements": [{"tag": "i", "id": "", "elements": []}, {"id": "", "tag": "#text", "text": "Contact"}]}]},
                    {"tag": "li", "id": "", "elements": [{"tag": "a", "id": "", "elements": [{"tag": "i", "id": "", "elements": []}, {"id": "", "tag": "#text", "text": "About"}]}]},
                ],
            }
        ],
    })

    data = {"data": [chunk_str]}
    response = requests.post(f"{url}/analyze", json=data)

    if response.status_code == 200:
        try:
            result = response.json()
            print("Valid JSON Response:", json.dumps(result, indent=2))
        except json.JSONDecodeError:
            print("Not valid JSON:", response.text)
            raise AssertionError("Invalid JSON response")

        output = result["output"]
        assert isinstance(output, dict), "Output is not a dictionary"
        assert "signal" in output, "Missing 'signal' key"
        assert "summary" in output, "Missing 'summary' key"
        # assert "actions" in output, "Missing 'actions' key"
        # assert "sections" in output, "Missing 'sections' key"
        assert output["signal"] in ["a1", "a0", "e0", "e1"], "Invalid 'signal' value"
    else:
        print(f"Request failed with status code: {response.status_code}")
        raise HTTPException(status_code=response.status_code, detail="Request to /analyze failed")

    print("Analyze endpoint test passed.")

def test_process_chunks():
    chunk_str = json.dumps({
        "tag": "nav",
        "id": "",
        "elements": [
            {
                "tag": "ul",
                "id": "",
                "elements": [
                    {"tag": "li", "id": "", "elements": [{"tag": "a", "id": "", "elements": [{"tag": "i", "id": "", "elements": []}, {"id": "", "tag": "#text", "text": "Category"}]}]},
                    {"tag": "li", "id": "", "elements": [{"tag": "a", "id": "", "elements": [{"tag": "i", "id": "", "elements": []}, {"id": "", "tag": "#text", "text": "Articles"}]}]},
                    {"tag": "li", "id": "", "elements": [{"tag": "a", "id": "", "elements": [{"tag": "i", "id": "", "elements": []}, {"id": "", "tag": "#text", "text": "Contact"}]}]},
                    {"tag": "li", "id": "", "elements": [{"tag": "a", "id": "", "elements": [{"tag": "i", "id": "", "elements": []}, {"id": "", "tag": "#text", "text": "About"}]}]},
                ],
            }
        ],
    })

    chunks = [chunk_str]
    msg = analysis_system_message
    processed = process_chunks(chunks, msg, limit=512)

    for chunk in processed:
        assert isinstance(chunk, str), f"chunk is not a string {chunk}"
        assert isinstance(json.loads(chunk), dict), "chunk is not valid JSON"
        assert compute_input_size(json.dumps(chunk), msg) <= 512, "chunk exceeds size limit"

    print("process_chunks function test passed.")