From 0bc7bdd550448fd82fd3914d1a4347350f5fff5c Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 14 Nov 2024 00:02:15 +0530
Subject: [PATCH] fix: fix test

---
 tests/logging_callback_tests/test_unit_tests_init_callbacks.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py
index ebc7dd33c..38883fa38 100644
--- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py
+++ b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py
@@ -31,6 +31,7 @@ from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
 from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
 from litellm.integrations.opik.opik import OpikLogger
 from litellm.integrations.opentelemetry import OpenTelemetry
+from litellm.integrations.mlflow import MlflowLogger
 from litellm.integrations.argilla import ArgillaLogger
 from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler
 from unittest.mock import patch
@@ -59,6 +60,7 @@ callback_class_str_to_classType = {
     "logfire": OpenTelemetry,
     "arize": OpenTelemetry,
     "langtrace": OpenTelemetry,
+    "mlflow": MlflowLogger,
 }
 
 expected_env_vars = {