from __future__ import annotations

import pytest

starlette = pytest.importorskip("starlette")

from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Route
from starlette.testclient import TestClient

from a2a_llm_tracker.middleware import TrackerMiddleware


def test_tracker_middleware_preserves_existing_response_trace_id():
    async def endpoint(request):
        return JSONResponse({"ok": True}, headers={"X-Trace-ID": "trace-downstream"})

    app = Starlette(routes=[Route("/", endpoint)])
    app.add_middleware(
        TrackerMiddleware,
        generate_request_id=False,
        generate_session_id=False,
        generate_call_id=False,
    )

    with TestClient(app) as client:
        response = client.get("/", headers={"X-Trace-ID": "trace-client"})

    assert response.status_code == 200
    assert response.headers["x-trace-id"] == "trace-downstream"
