@@ -59,8 +59,8 @@ async def test_initialize_cosmos(cosmos_db_client, mocker):
5959 mock_file_container = mock .MagicMock ()
6060 mock_log_container = mock .MagicMock ()
6161
62- # Use AsyncMock to mock asynchronous container creation
63- mock_database .create_container = AsyncMock (side_effect = [
62+ # Mock get_container_client method (since _get_container uses this)
63+ mock_database .get_container_client = mock . MagicMock (side_effect = [
6464 mock_batch_container ,
6565 mock_file_container ,
6666 mock_log_container
@@ -69,10 +69,10 @@ async def test_initialize_cosmos(cosmos_db_client, mocker):
6969 # Call the initialize_cosmos method
7070 await cosmos_db_client .initialize_cosmos ()
7171
72- # Assert that the containers were created or fetched successfully
73- mock_database .create_container .assert_any_call (id = batch_container , partition_key = mock . ANY )
74- mock_database .create_container .assert_any_call (id = file_container , partition_key = mock . ANY )
75- mock_database .create_container .assert_any_call (id = log_container , partition_key = mock . ANY )
72+ # Assert that the containers were fetched successfully
73+ mock_database .get_container_client .assert_any_call (batch_container )
74+ mock_database .get_container_client .assert_any_call (file_container )
75+ mock_database .get_container_client .assert_any_call (log_container )
7676
7777 # Check the client and containers were set
7878 assert cosmos_db_client .client is not None
@@ -87,15 +87,15 @@ async def test_initialize_cosmos_with_error(cosmos_db_client, mocker):
8787 mock_client = mocker .patch .object (CosmosClient , 'get_database_client' , return_value = mock .MagicMock ())
8888 mock_database = mock_client .return_value
8989
90- # Simulate a general exception during container creation
91- mock_database .create_container = AsyncMock (side_effect = Exception ("Failed to create container" ))
90+ # Simulate a general exception during container access
91+ mock_database .get_container_client = mock . MagicMock (side_effect = Exception ("Failed to get container" ))
9292
9393 # Call the initialize_cosmos method and expect it to raise an error
9494 with pytest .raises (Exception ) as exc_info :
9595 await cosmos_db_client .initialize_cosmos ()
9696
9797 # Assert that the exception message matches the expected message
98- assert str (exc_info .value ) == "Failed to create container"
98+ assert str (exc_info .value ) == "Failed to get container"
9999
100100
101101@pytest .mark .asyncio
@@ -104,16 +104,13 @@ async def test_initialize_cosmos_container_exists_error(cosmos_db_client, mocker
104104 mock_client = mocker .patch .object (CosmosClient , 'get_database_client' , return_value = mock .MagicMock ())
105105 mock_database = mock_client .return_value
106106
107- # Simulating CosmosResourceExistsError for container creation
108- mock_database .create_container = AsyncMock (side_effect = CosmosResourceExistsError )
109-
110107 # Use AsyncMock for asynchronous methods
111108 mock_batch_container = mock .MagicMock ()
112109 mock_file_container = mock .MagicMock ()
113110 mock_log_container = mock .MagicMock ()
114111
115- # Use AsyncMock to mock asynchronous container creation
116- mock_database .create_container = AsyncMock (side_effect = [
112+ # Mock get_container_client method to return existing containers
113+ mock_database .get_container_client = mock . MagicMock (side_effect = [
117114 mock_batch_container ,
118115 mock_file_container ,
119116 mock_log_container
@@ -122,10 +119,10 @@ async def test_initialize_cosmos_container_exists_error(cosmos_db_client, mocker
122119 # Call the initialize_cosmos method
123120 await cosmos_db_client .initialize_cosmos ()
124121
125- # Assert that the container creation method was called with the correct arguments
126- mock_database .create_container .assert_any_call (id = 'batch_container' , partition_key = mock . ANY )
127- mock_database .create_container .assert_any_call (id = 'file_container' , partition_key = mock . ANY )
128- mock_database .create_container .assert_any_call (id = 'log_container' , partition_key = mock . ANY )
122+ # Assert that the container access method was called with the correct arguments
123+ mock_database .get_container_client .assert_any_call ('batch_container' )
124+ mock_database .get_container_client .assert_any_call ('file_container' )
125+ mock_database .get_container_client .assert_any_call ('log_container' )
129126
130127 # Check that existing containers are returned (mocked containers)
131128 assert cosmos_db_client .batch_container == mock_batch_container
0 commit comments