fix(tests): drain SDK logging before restoring callbacks (#46924) (#47143)

(cherry picked from commit 1451aed726)

Co-authored-by: Kerry Lu <klu@berri.ai>
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/integration/database/test_harness_lifecycle.py

@@ -106,6 +106,82 @@ def test_runner_environment_drops_inherited_provider_api_keys() -> None:
     assert inherited["OPENAI_API_KEY"] == "ambient-openai"
 
 
+@pytest.mark.parametrize("fails", (False, True), ids=("passing-test", "failing-test"))
+def test_sdk_teardown_delivers_logging_before_restoring_callbacks(tmp_path: Path, fails: bool) -> None:
+    root: Final = Path(__file__).resolve().parents[3]
+    events: Final = tmp_path / "callbacks.json"
+    probe: Final = tmp_path / "test_callback_teardown.py"
+    probe.write_text(
+        "import asyncio, json, litellm, pytest\n"
+        "from pathlib import Path\n"
+        "from litellm.integrations.custom_logger import CustomLogger\n"
+        f"EVENTS = Path({str(events)!r})\n"
+        "ORIGINAL = tuple(litellm.callbacks)\n"
+        "@pytest.fixture\n"
+        "def release():\n"
+        "    gate = asyncio.Event()\n"
+        "    yield gate\n"
+        "    gate.set()\n"
+        "async def test_request(release):\n"
+        "    started = asyncio.Event()\n"
+        "    request_loop = asyncio.get_running_loop()\n"
+        "    class Recorder(CustomLogger):\n"
+        "        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):\n"
+        "            started.set()\n"
+        "            await release.wait()\n"
+        "            EVENTS.write_text(json.dumps({\n"
+        "                'registered': self in litellm.callbacks,\n"
+        "                'same_loop': asyncio.get_running_loop() is request_loop,\n"
+        "                'call_id': kwargs['litellm_call_id'],\n"
+        "            }))\n"
+        "    litellm.callbacks.append(Recorder())\n"
+        "    response = await litellm.acompletion(model='gpt-4o-mini',\n"
+        "        messages=[{'role': 'user', 'content': 'teardown'}],\n"
+        "        mock_response='done', litellm_call_id='teardown-request')\n"
+        "    await asyncio.wait_for(started.wait(), timeout=5)\n"
+        "    assert response.choices[0].message.content == 'done'\n"
+        "    assert not EVENTS.exists()\n"
+        f"    assert not {fails}, 'intentional test failure'\n"
+        "def test_next():\n"
+        "    assert EVENTS.exists(), 'logging did not finish before the next test'\n"
+        "    assert json.loads(EVENTS.read_text()) == {\n"
+        "        'registered': True, 'same_loop': True, 'call_id': 'teardown-request'}\n"
+        "    assert tuple(litellm.callbacks) == ORIGINAL\n"
+    )
+    result: Final = subprocess.run(
+        (
+            sys.executable,
+            "-P",
+            "-m",
+            "pytest",
+            str(probe),
+            "--noconftest",
+            "-p",
+            "tests.integration.sdk.conftest",
+            "-o",
+            "addopts=",
+            "-o",
+            "asyncio_mode=auto",
+            "-o",
+            "asyncio_default_fixture_loop_scope=session",
+            f"--junitxml={tmp_path / 'junit.xml'}",
+            "-q",
+        ),
+        cwd=root,
+        env=_nested_pytest_environment(os.environ),
+        capture_output=True,
+        text=True,
+        timeout=60,
+        check=False,
+    )
+    assert result.returncode == (1 if fails else 0), result.stdout + result.stderr
+    cases: Final = ElementTree.parse(tmp_path / "junit.xml").findall(".//testcase")
+    assert tuple(case.attrib["name"] for case in cases) == ("test_request", "test_next")
+    assert all(case.find("error") is None for case in cases), result.stdout + result.stderr
+    assert (cases[0].find("failure") is not None) == fails, result.stdout + result.stderr
+    assert cases[1].find("failure") is None, result.stdout + result.stderr
+
+
 def test_owned_readiness_rejects_a_healthy_listener_from_another_process(
     monkeypatch: pytest.MonkeyPatch, tmp_path: Path
 ) -> None:

tests/integration/sdk/conftest.py

@@ -1,7 +1,11 @@
+from collections.abc import AsyncIterator
 from typing import Final
 
-import litellm
 import pytest
+import pytest_asyncio
+
+import litellm
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 
 CALLBACK_LISTS: Final = (
     "callbacks",
@@ -14,7 +18,12 @@ CALLBACK_LISTS: Final = (
 )
 
 
-@pytest.fixture(autouse=True)
-def isolate_litellm_callback_lists(monkeypatch: pytest.MonkeyPatch) -> None:
+@pytest_asyncio.fixture(autouse=True, loop_scope="function")
+async def isolate_litellm_callback_lists(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[None]:
     for name in CALLBACK_LISTS:
         monkeypatch.setattr(litellm, name, list(getattr(litellm, name)))
+    try:
+        yield
+    finally:
+        await GLOBAL_LOGGING_WORKER.flush()
+        await GLOBAL_LOGGING_WORKER.stop()