From 9416bc0535753e39142394ac13844ca4d020d2cd Mon Sep 17 00:00:00 2001
From: Tvrtko Sternak <sternakt@gmail.com>
Date: Tue, 21 Jan 2025 14:00:39 +0100
Subject: [PATCH] Cleanup conftest

---
 test/conftest.py | 45 +++++++++++++++++++++++----------------------
 1 file changed, 23 insertions(+), 22 deletions(-)

diff --git a/test/conftest.py b/test/conftest.py
index 1b8b0e793..7afee1462 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -10,6 +10,7 @@
 from typing import Any, Optional
 
 import pytest
+from pytest import CallInfo, Item
 
 import autogen
 
@@ -119,6 +120,28 @@ def api_key(self) -> str:
         return self.llm_config["config_list"][0]["api_key"]  # type: ignore[no-any-return]
 
 
+class CensoredError(Exception):
+    def __init__(self, exception: BaseException):
+        self.exception = exception
+        self.__traceback__ = exception.__traceback__
+        original_message = "".join([repr(arg) for arg in exception.args])
+        message = Secrets.sanitize_secrets(original_message)
+        super().__init__(message)
+
+
+def pytest_runtest_makereport(item: Item, call: CallInfo[Any]) -> None:
+    """
+    Hook to customize the exception output.
+    This is called after each test call.
+    """
+    if call.excinfo is not None:  # This means the test failed
+        original_message = "".join([repr(arg) for arg in call.excinfo.value.args])
+
+        if Secrets.needs_sanitizing(original_message):
+            censored_exception = CensoredError(call.excinfo.value)
+            call.excinfo = pytest.ExceptionInfo.from_exception(censored_exception)
+
+
 def get_credentials(
     filter_dict: Optional[dict[str, Any]] = None, temperature: float = 0.0, fail_if_empty: bool = True
 ) -> Credentials:
@@ -286,28 +309,6 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
         session.exitstatus = 0
 
 
-class CensoredError(Exception):
-    def __init__(self, exception: Exception):
-        self.exception = exception
-        self.__traceback__ = exception.__traceback__
-        original_message = "".join([repr(arg) for arg in exception.args])
-        message = Secrets.sanitize_secrets(original_message)
-        super().__init__(message)
-
-
-# def pytest_runtest_makereport(item, call):
-#     """
-#     Hook to customize the exception output.
-#     This is called after each test call.
-#     """
-#     if call.excinfo is not None:  # This means the test failed
-#         original_message = "".join([repr(arg) for arg in call.excinfo.value.args])
-
-#         if Secrets.needs_sanitizing(original_message):
-#             censored_exception = CensoredError(call.excinfo.value)
-#             call.excinfo = pytest.ExceptionInfo.from_exception(censored_exception)
-
-
 credentials_all_llms = [
     pytest.param(
         credentials_gpt_4o_mini.__name__,