from sanic.server import HttpProtocol
from sanic.response import text


class CustomHttpProtocol(HttpProtocol):
    def write_response(self, response):
        if isinstance(response, str):
            response = text(response)
        self.transport.write(response.output(self.request.version))
        self.transport.close()


def test_use_custom_protocol(app):
    @app.route("/1")
    async def handler_1(request):
        return "OK"

    server_kwargs = {"protocol": CustomHttpProtocol}
    request, response = app.test_client.get("/1", server_kwargs=server_kwargs)
    assert response.status == 200
    assert response.text == "OK"