diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py
index ac4a1a11a80..d3b04c8bb6f 100644
--- a/tests/test_litellm_rust/ocr/test_callbacks.py
+++ b/tests/test_litellm_rust/ocr/test_callbacks.py
@@ -12,6 +12,7 @@ from hypothesis import HealthCheck, given, settings
 from hypothesis import strategies as st
 
 import litellm
+from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
 from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
@@ -420,7 +421,7 @@ async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_f
     class Blocked(Exception):
         pass
 
-    class Block(CustomLogger):
+    class Block(CustomGuardrail):
         async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
             request_data["litellm_logging_obj"].model_call_details["blocked-by"] = token
             raise Blocked("blocked after the provider answered")