diff --git a/api/tests/unit_tests/core/helper/test_encrypter.py b/api/tests/unit_tests/core/helper/test_encrypter.py
index 61cf8f255d..00749f7c72 100644
--- a/api/tests/unit_tests/core/helper/test_encrypter.py
+++ b/api/tests/unit_tests/core/helper/test_encrypter.py
@@ -45,6 +45,7 @@ class TestEncryptToken:
         mock_tenant = MagicMock()
         mock_tenant.encrypt_public_key = "mock_public_key"
         mock_query.return_value.filter.return_value.first.return_value = mock_tenant
+        mock_query.return_value.where.return_value.first.return_value = mock_tenant
         mock_encrypt.return_value = b"encrypted_data"
 
         result = encrypt_token("tenant-123", "test_token")
diff --git a/api/tests/unit_tests/factories/test_build_from_mapping.py b/api/tests/unit_tests/factories/test_build_from_mapping.py
index 48463a369e..d87b7f5327 100644
--- a/api/tests/unit_tests/factories/test_build_from_mapping.py
+++ b/api/tests/unit_tests/factories/test_build_from_mapping.py
@@ -54,8 +54,12 @@ def mock_tool_file():
     mock.mimetype = "application/pdf"
     mock.original_url = "http://example.com/tool.pdf"
     mock.size = 2048
-    with patch("factories.file_factory.db.session.query") as mock_query:
+    with (
+        patch("factories.file_factory.db.session.query") as mock_query,
+        patch("factories.file_factory.db.session.scalar") as mock_query1,
+    ):
         mock_query.return_value.filter.return_value.first.return_value = mock
+        mock_query1.return_value.filter.return_value.first.return_value = mock
         yield mock
 
 
diff --git a/api/tests/unit_tests/services/auth/test_api_key_auth_service.py b/api/tests/unit_tests/services/auth/test_api_key_auth_service.py
index f0e425e742..e8ddb57e44 100644
--- a/api/tests/unit_tests/services/auth/test_api_key_auth_service.py
+++ b/api/tests/unit_tests/services/auth/test_api_key_auth_service.py
@@ -139,6 +139,7 @@ class TestApiKeyAuthService:
         mock_binding = Mock()
         mock_binding.credentials = json.dumps(self.mock_credentials)
         mock_session.query.return_value.filter.return_value.first.return_value = mock_binding
+        mock_session.query.return_value.where.return_value.first.return_value = mock_binding
 
         result = ApiKeyAuthService.get_auth_credentials(self.tenant_id, self.category, self.provider)