Browser_Agent/tests/test_api_server.py

import pytest
from fastapi.testclient import TestClient
from api.server import app
from job_queue.task_queue import task_queue
from workflows.ephemeral_media import ephemeral_media_store

client = TestClient(app)

def test_health_endpoint():
    response = client.get("/health")
    assert response.status_code == 200
    assert response.json()["status"] == "ok"
    assert "queue_size" in response.json()

def test_create_task():
    response = client.post("/task", json={
        "task_type": "test_task",
        "params": {"a": 1}
    })
    assert response.status_code == 200
    assert "task_id" in response.json()
    task_id = response.json()["task_id"]
    
    # Check status
    status_res = client.get(f"/task/{task_id}")
    assert status_res.status_code == 200
    assert status_res.json()["status"] == "pending"

def test_task_not_found():
    res = client.get("/task/invalid_task_id")
    assert res.status_code == 404

def test_ephemeral_media_endpoints():
    media_id = ephemeral_media_store.put(b"dummy_data", "image/png", "test.png")
    
    # Get media
    res = client.get(f"/api/media/{media_id}")
    assert res.status_code == 200
    assert res.content == b"dummy_data" # assuming it bypasses sanitizer due to invalid image or just returns it
    assert "no-store" in res.headers.get("Cache-Control", "")
    
    # Get encrypted media
    res_enc = client.get(f"/api/media/{media_id}/encrypted")
    assert res_enc.status_code == 200
    assert res_enc.content != b"dummy_data"
    
    # Delete media
    res_del = client.delete(f"/api/media/{media_id}")
    assert res_del.status_code == 200
    
    # Get again -> 404
    res_404 = client.get(f"/api/media/{media_id}")
    assert res_404.status_code == 404

def test_downloads_path_traversal():
    res = client.get("/downloads/..%2Fetc%2Fpasswd")
    assert res.status_code in (403, 404)