fix: fix unit tests import

Author: JackYPCOnline

tests/strands/session/test_repository_session_manager.py

@@ -1,5 +1,7 @@
 """Tests for AgentSessionManager."""
 
+from unittest.mock import Mock
+
 import pytest
 
 from strands.agent.agent import Agent
@@ -34,7 +36,6 @@ def agent():
 @pytest.fixture
 def mock_multi_agent():
     """Create mock multi-agent for testing."""
-    from unittest.mock import Mock
 
     mock = Mock()
     mock.id = "test-multi-agent"
@@ -240,7 +241,6 @@ def test_initialize_multi_agent_existing(multi_agent_session_manager, mock_multi
     multi_agent_session_manager.session_repository.create_multi_agent("test-multi-session", mock_multi_agent)
 
     # Create a mock with updated state for the update call
-    from unittest.mock import Mock
 
     updated_mock = Mock()
     updated_mock.id = "test-multi-agent"

tests/strands/session/test_s3_session_manager.py

@@ -1,6 +1,7 @@
 """Tests for S3SessionManager."""
 
 import json
+from unittest.mock import Mock
 
 import boto3
 import pytest
@@ -379,7 +380,6 @@ def test__get_message_path_invalid_message_id(message_id, s3_manager):
 @pytest.fixture
 def mock_multi_agent():
     """Create mock multi-agent for testing."""
-    from unittest.mock import Mock
 
     mock = Mock()
     mock.id = "test-multi-agent"
@@ -452,9 +452,6 @@ def test_update_multi_agent(multi_agent_manager, multi_agent_session, mock_multi
     multi_agent_manager.create_session(multi_agent_session)
     multi_agent_manager.create_multi_agent(multi_agent_session.session_id, mock_multi_agent)
 
-    # Update multi-agent - create a new mock with updated state
-    from unittest.mock import Mock
-
     updated_mock = Mock()
     updated_mock.id = mock_multi_agent.id
     updated_mock.serialize_state.return_value = {"id": mock_multi_agent.id, "state": {"updated": "value"}}
@@ -470,9 +467,6 @@ def test_update_nonexistent_multi_agent(multi_agent_manager, multi_agent_session
     # Create session
     multi_agent_manager.create_session(multi_agent_session)
 
-    # Update nonexistent multi-agent
-    from unittest.mock import Mock
-
     nonexistent_mock = Mock()
     nonexistent_mock.id = "nonexistent"
     with pytest.raises(SessionException):