From 3d0579c7eaa2d7ce51cee2da80e37762aef3e9fe Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 23 Oct 2024 10:06:43 +0530
Subject: [PATCH] add unit testing for langfuse logging

---
 .../test_langfuse_unit_tests.py               | 108 ++++++++++++++++++
 1 file changed, 108 insertions(+)

diff --git a/tests/logging_callback_tests/test_langfuse_unit_tests.py b/tests/logging_callback_tests/test_langfuse_unit_tests.py
index cacffe7e6a6..2a6cbe00a90 100644
--- a/tests/logging_callback_tests/test_langfuse_unit_tests.py
+++ b/tests/logging_callback_tests/test_langfuse_unit_tests.py
@@ -20,6 +20,7 @@ from litellm.integrations.langfuse.langfuse import (
 from litellm.integrations.langfuse.langfuse_handler import LangFuseHandler
 from litellm.types.utils import StandardCallbackDynamicParams
 from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache
+from unittest.mock import Mock, patch
 
 
 @pytest.fixture
@@ -109,3 +110,110 @@ def test_get_langfuse_logger_for_request_with_no_dynamic_params(
         assert result.public_key == "global_public_key"
         assert result.secret_key == "global_secret"
         assert result.langfuse_host == "https://global.langfuse.com"
+
+
+def test_dynamic_langfuse_credentials_are_passed():
+    # Test when credentials are passed
+    params_with_credentials = StandardCallbackDynamicParams(
+        langfuse_public_key="test_key",
+        langfuse_secret="test_secret",
+        langfuse_host="https://test.langfuse.com",
+    )
+    assert (
+        LangFuseHandler._dynamic_langfuse_credentials_are_passed(
+            params_with_credentials
+        )
+        is True
+    )
+
+    # Test when no credentials are passed
+    params_without_credentials = StandardCallbackDynamicParams()
+    assert (
+        LangFuseHandler._dynamic_langfuse_credentials_are_passed(
+            params_without_credentials
+        )
+        is False
+    )
+
+    # Test when only some credentials are passed
+    params_partial_credentials = StandardCallbackDynamicParams(
+        langfuse_public_key="test_key"
+    )
+    assert (
+        LangFuseHandler._dynamic_langfuse_credentials_are_passed(
+            params_partial_credentials
+        )
+        is True
+    )
+
+
+def test_get_dynamic_langfuse_logging_config():
+    # Test with dynamic params
+    dynamic_params = StandardCallbackDynamicParams(
+        langfuse_public_key="dynamic_key",
+        langfuse_secret="dynamic_secret",
+        langfuse_host="https://dynamic.langfuse.com",
+    )
+    config = LangFuseHandler.get_dynamic_langfuse_logging_config(dynamic_params)
+    assert config["langfuse_public_key"] == "dynamic_key"
+    assert config["langfuse_secret"] == "dynamic_secret"
+    assert config["langfuse_host"] == "https://dynamic.langfuse.com"
+
+    # Test with no dynamic params
+    empty_params = StandardCallbackDynamicParams()
+    config = LangFuseHandler.get_dynamic_langfuse_logging_config(empty_params)
+    assert config["langfuse_public_key"] is None
+    assert config["langfuse_secret"] is None
+    assert config["langfuse_host"] is None
+
+
+def test_return_global_langfuse_logger():
+    mock_cache = Mock()
+    global_logger = LangFuseLogger(
+        langfuse_public_key="global_key", langfuse_secret="global_secret"
+    )
+
+    # Test with existing global logger
+    result = LangFuseHandler._return_global_langfuse_logger(global_logger, mock_cache)
+    assert result == global_logger
+
+    # Test without global logger, but with cached logger, should return cached logger
+    mock_cache.get_cache.return_value = global_logger
+    result = LangFuseHandler._return_global_langfuse_logger(None, mock_cache)
+    assert result == global_logger
+
+    # Test without global logger and without cached logger, should create new logger
+    mock_cache.get_cache.return_value = None
+    with patch.object(
+        LangFuseHandler,
+        "_create_langfuse_logger_from_credentials",
+        return_value=global_logger,
+    ):
+        result = LangFuseHandler._return_global_langfuse_logger(None, mock_cache)
+        assert result == global_logger
+
+
+def test_get_langfuse_logger_for_request_with_cached_logger():
+    """
+    Test that get_langfuse_logger_for_request returns the cached logger if it exists when dynamic params are passed
+    """
+    mock_cache = Mock()
+    cached_logger = LangFuseLogger(
+        langfuse_public_key="cached_key", langfuse_secret="cached_secret"
+    )
+    mock_cache.get_cache.return_value = cached_logger
+
+    dynamic_params = StandardCallbackDynamicParams(
+        langfuse_public_key="test_key",
+        langfuse_secret="test_secret",
+        langfuse_host="https://test.langfuse.com",
+    )
+
+    result = LangFuseHandler.get_langfuse_logger_for_request(
+        standard_callback_dynamic_params=dynamic_params,
+        in_memory_dynamic_logger_cache=mock_cache,
+        globalLangfuseLogger=None,
+    )
+
+    assert result == cached_logger
+    mock_cache.get_cache.assert_called_once()