asphalt-framework / asphalt-web Goto Github PK

from contextlib import asynccontextmanager

from fastapi import FastAPI


def fake_answer_to_everything_ml_model(x: float):
    return x * 42


ml_models = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load the ML model
    ml_models["answer_to_everything"] = fake_answer_to_everything_ml_model
    yield
    # Clean up the ML models and release the resources
    ml_models.clear()


app = FastAPI(lifespan=lifespan)


@app.get("/predict")
async def predict(x: float):
    result = ml_models["answer_to_everything"](x)
    return {"result": result}

from asphalt.core import Context
from asphalt.web.fastapi import FastAPIComponent
from fastapi import FastAPI
from websockets import connect

import pytest

@pytest.mark.asyncio
async def test(unused_tcp_port):
    application = FastAPI()

    async with Context() as ctx:
        await FastAPIComponent(
            components={},
            app=application,
            port=unused_tcp_port,
        ).start(ctx)

        async with connect(f"ws://127.0.0.1:{unused_tcp_port}/wrong_ws"):
            pass
        # websockets.exceptions.InvalidStatusCode: server rejected WebSocket connection: HTTP 403