diff --git a/tests/test_dependency_contextmanager.py b/tests/test_dependency_contextmanager.py index 02c10458c..33e204452 100644 --- a/tests/test_dependency_contextmanager.py +++ b/tests/test_dependency_contextmanager.py @@ -16,6 +16,12 @@ state = { "context_b": "not started b", "bg": "not set", "sync_bg": "not set", + "sync_bg_task_pre": "not executed", + "sync_bg_task_post": "not executed", + "async_bg_task_pre": "not executed", + "async_bg_task_post": "not executed", + "sync_bg_endpoint": "not started", + "async_bg_endpoint": "not started", } errors = [] @@ -87,6 +93,58 @@ async def context_b(state: dict = Depends(context_a)): state["context_b"] = f"finished b with a: {state['context_a']}" +def sync_dependency_with_bg_tasks( + background_tasks: BackgroundTasks, + state: dict = Depends(get_state), +): + """Sync generator dependency that adds background tasks before and after yield.""" + + def sync_bg_task_pre(): + state["sync_bg_task_pre"] = "executed" + + async def async_bg_task_pre(): + state["async_bg_task_pre"] = "executed" + + def sync_bg_task_post(): + state["sync_bg_task_post"] = "executed" + + async def async_bg_task_post(): + state["async_bg_task_post"] = "executed" + + state["sync_bg_endpoint"] = "started" + background_tasks.add_task(sync_bg_task_pre) + background_tasks.add_task(async_bg_task_pre) + yield state + background_tasks.add_task(sync_bg_task_post) + background_tasks.add_task(async_bg_task_post) + + +async def async_dependency_with_bg_tasks( + background_tasks: BackgroundTasks, + state: dict = Depends(get_state), +): + """Async generator dependency that adds background tasks before and after yield.""" + + def sync_bg_task_pre(): + state["sync_bg_task_pre"] = "executed" + + async def async_bg_task_pre(): + state["async_bg_task_pre"] = "executed" + + def sync_bg_task_post(): + state["sync_bg_task_post"] = "executed" + + async def async_bg_task_post(): + state["async_bg_task_post"] = "executed" + + state["async_bg_endpoint"] = "started" + background_tasks.add_task(sync_bg_task_pre) + background_tasks.add_task(async_bg_task_pre) + yield state + background_tasks.add_task(sync_bg_task_post) + background_tasks.add_task(async_bg_task_post) + + @app.get("/async") async def get_async(state: str = Depends(asyncgen_state)): return state @@ -142,6 +200,20 @@ async def get_context_b_bg(tasks: BackgroundTasks, state: dict = Depends(context return state +@app.get("/sync_bg_after_yield") +def sync_bg_after_yield_endpoint( + state: dict = Depends(sync_dependency_with_bg_tasks), +): + return state + + +@app.get("/async_bg_after_yield") +async def async_bg_after_yield_endpoint( + state: dict = Depends(async_dependency_with_bg_tasks), +): + return state + + # Sync versions @@ -398,3 +470,47 @@ def test_sync_background_tasks(): assert state["context_b"] == "finished b with a: started a" assert state["context_a"] == "finished a" assert state["sync_bg"] == "sync_bg set - b: started b - a: started a" + + +def test_sync_dependency_background_tasks_after_yield(): + # Reset states before test + global state + state["sync_bg_task_pre"] = "not executed" + state["async_bg_task_pre"] = "not executed" + state["sync_bg_task_post"] = "not executed" + state["async_bg_task_post"] = "not executed" + + response = client.get("/sync_bg_after_yield") + assert response.status_code == 200 + data = response.json() + assert data["sync_bg_endpoint"] == "started" + assert data["sync_bg_task_pre"] == "not executed" + assert data["async_bg_task_pre"] == "not executed" + assert data["sync_bg_task_post"] == "not executed" + assert data["async_bg_task_post"] == "not executed" + assert state["sync_bg_task_pre"] == "executed" + assert state["async_bg_task_pre"] == "executed" + assert state["sync_bg_task_post"] == "executed" + assert state["async_bg_task_post"] == "executed" + + +def test_async_dependency_background_tasks_after_yield(): + # Reset states before test + global state + state["sync_bg_task_pre"] = "not executed" + state["async_bg_task_pre"] = "not executed" + state["sync_bg_task_post"] = "not executed" + state["async_bg_task_post"] = "not executed" + + response = client.get("/async_bg_after_yield") + assert response.status_code == 200 + data = response.json() + assert data["async_bg_endpoint"] == "started" + assert data["sync_bg_task_pre"] == "not executed" + assert data["async_bg_task_pre"] == "not executed" + assert data["sync_bg_task_post"] == "not executed" + assert data["async_bg_task_post"] == "not executed" + assert state["sync_bg_task_pre"] == "executed" + assert state["async_bg_task_pre"] == "executed" + assert state["sync_bg_task_post"] == "executed" + assert state["async_bg_task_post"] == "executed"