import pytest
import requests
from fastapi import FastAPI
import starlette

import ray
from ray import serve
from ray.serve._private.constants import RAY_SERVE_REQUEST_ID_HEADER


def test_request_id_header_by_default(serve_instance):
    """Test that a request_id is generated by default and returned as a header."""

    @serve.deployment
    class Model:
        def __call__(self):
            request_id = ray.serve.context._serve_request_context.get().request_id
            return request_id

    serve.run(Model.bind())
    resp = requests.get("http://localhost:8000")
    assert resp.status_code == 200
    assert RAY_SERVE_REQUEST_ID_HEADER in resp.headers
    assert resp.text == resp.headers[RAY_SERVE_REQUEST_ID_HEADER]


class TestUserProvidedRequestIDHeader:
    def verify_result(self):
        resp = requests.get(
            "http://localhost:8000", headers={RAY_SERVE_REQUEST_ID_HEADER: "123-234"}
        )
        assert resp.status_code == 200
        assert resp.json() == 1
        assert RAY_SERVE_REQUEST_ID_HEADER in resp.headers
        assert resp.headers[RAY_SERVE_REQUEST_ID_HEADER] == "123-234"

    def test_basic(self, serve_instance):
        @serve.deployment
        class Model:
            def __call__(self) -> int:
                request_id = ray.serve.context._serve_request_context.get().request_id
                assert request_id == "123-234"
                return 1

        serve.run(Model.bind())
        self.verify_result()

    def test_fastapi(self, serve_instance):
        app = FastAPI()

        @serve.deployment
        @serve.ingress(app)
        class Model:
            @app.get("/")
            def say_hi(self) -> int:
                request_id = ray.serve.context._serve_request_context.get().request_id
                assert request_id == "123-234"
                return 1

        serve.run(Model.bind())
        self.verify_result()

    def test_starlette_resp(self, serve_instance):
        @serve.deployment
        class Model:
            def __call__(self) -> int:
                request_id = ray.serve.context._serve_request_context.get().request_id
                assert request_id == "123-234"
                return starlette.responses.Response("1", media_type="application/json")

        serve.run(Model.bind())
        self.verify_result()


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", "-s", __file__]))
