Fix the unit tests

microsoft · Jun 19, 2024 · e5174c1 · e5174c1
1 parent 4aa135f
commit e5174c1
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/promptflow-azure/promptflow/azure/_utils/general.py b/src/promptflow-azure/promptflow/azure/_utils/general.py
@@ -4,10 +4,9 @@
 
 import jwt
 
+from promptflow.azure._utils._token_cache import ArmTokenCache
 from promptflow.core._connection_provider._utils import get_arm_token, get_token
 
-from ._token_cache import ArmTokenCache
-
 
 def is_arm_id(obj) -> bool:
     return isinstance(obj, str) and obj.startswith("azureml://")

diff --git a/src/promptflow-evals/tests/evals/unittests/test_eval_run.py b/src/promptflow-evals/tests/evals/unittests/test_eval_run.py
@@ -4,6 +4,7 @@
 from unittest.mock import MagicMock, patch
 from uuid import uuid4
 
+import jwt
 import pytest
 
 import promptflow.evals.evaluate._utils as ev_utils
@@ -18,8 +19,13 @@ def setup_data():
     Singleton._instances.clear()
 
 
+def generate_mock_token():
+    expiration_time = time.time() + 3600  # 1 hour in the future
+    return jwt.encode({"exp": expiration_time}, "secret", algorithm="HS256")
+
+
 @pytest.mark.unittest
-@patch.object(ArmTokenCache, "_fetch_token")
+@patch.object(ArmTokenCache, "_fetch_token", return_value=generate_mock_token())
 class TestEvalRun:
     """Unit tests for the eval-run object."""