1- # Copyright 2024 Google LLC
1+ # Copyright 2025 Google LLC
22#
33# Licensed under the Apache License, Version 2.0 (the "License");
44# you may not use this file except in compliance with the License.
@@ -421,6 +421,69 @@ async def test_aload_toolset_with_args(
421421 strict = True ,
422422 )
423423
424+ @pytest .mark .asyncio
425+ @patch ("toolbox_core.sync_client.ToolboxSyncClient.load_toolset" )
426+ async def test_aload_toolset_with_deprecated_args (
427+ self , mock_sync_core_load_toolset , toolbox_client
428+ ):
429+ mock_core_tool_instance = create_mock_core_sync_tool (
430+ model_name = "MyAsyncSetModel"
431+ )
432+ mock_sync_core_load_toolset .return_value = [mock_core_tool_instance ]
433+
434+ auth_tokens_deprecated = {"token_deprecated" : lambda : "value_dep" }
435+ auth_headers_deprecated = {"header_deprecated" : lambda : "value_head_dep" }
436+ bound_params = {"param1" : "value4" }
437+ toolset_name = "my_async_toolset"
438+
439+ # Scenario 2: auth_tokens and auth_headers provided, auth_token_getters is default (empty initially)
440+ with pytest .warns (DeprecationWarning ) as record :
441+ await toolbox_client .aload_toolset (
442+ toolset_name ,
443+ auth_tokens = auth_tokens_deprecated , # This will be used for auth_token_getters
444+ auth_headers = auth_headers_deprecated , # This will warn as auth_token_getters is now populated
445+ bound_params = bound_params ,
446+ )
447+ assert len (record ) == 2
448+ messages = sorted ([str (r .message ) for r in record ])
449+
450+ assert (
451+ messages [0 ]
452+ == "Argument `auth_tokens` is deprecated. Use `auth_token_getters` instead."
453+ )
454+ assert (
455+ messages [1 ]
456+ == "Both `auth_token_getters` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_token_getters` will be used."
457+ )
458+
459+ expected_getters_for_call = auth_tokens_deprecated
460+
461+ mock_sync_core_load_toolset .assert_called_with (
462+ name = toolset_name ,
463+ auth_token_getters = expected_getters_for_call ,
464+ bound_params = bound_params ,
465+ strict = False ,
466+ )
467+ mock_sync_core_load_toolset .reset_mock ()
468+
469+ with pytest .warns (
470+ DeprecationWarning ,
471+ match = "Argument `auth_headers` is deprecated. Use `auth_token_getters` instead." ,
472+ ) as record :
473+ await toolbox_client .aload_toolset (
474+ toolset_name ,
475+ auth_headers = auth_headers_deprecated ,
476+ bound_params = bound_params ,
477+ )
478+ assert len (record ) == 1
479+
480+ mock_sync_core_load_toolset .assert_called_with (
481+ name = toolset_name ,
482+ auth_token_getters = auth_headers_deprecated ,
483+ bound_params = bound_params ,
484+ strict = False ,
485+ )
486+
424487 @patch ("toolbox_langchain.client.ToolboxCoreSyncClient" )
425488 def test_init_with_client_headers (self , mock_core_client_constructor ):
426489 """Tests that client_headers are passed to the core client during initialization."""
@@ -429,3 +492,27 @@ def test_init_with_client_headers(self, mock_core_client_constructor):
429492 mock_core_client_constructor .assert_called_once_with (
430493 url = URL , client_headers = headers
431494 )
495+
496+ @patch ("toolbox_langchain.client.ToolboxCoreSyncClient" )
497+ def test_context_manager (self , mock_core_client_constructor ):
498+ """Tests that the client can be used as a context manager."""
499+ with ToolboxClient (URL ) as client :
500+ assert isinstance (client , ToolboxClient )
501+ mock_core_client_constructor .return_value .close .assert_not_called ()
502+ mock_core_client_constructor .return_value .close .assert_called_once ()
503+
504+ @pytest .mark .asyncio
505+ @patch ("toolbox_langchain.client.ToolboxCoreSyncClient" )
506+ async def test_async_context_manager (self , mock_core_client_constructor ):
507+ """Tests that the client can be used as an async context manager."""
508+ async with ToolboxClient (URL ) as client :
509+ assert isinstance (client , ToolboxClient )
510+ mock_core_client_constructor .return_value .close .assert_not_called ()
511+ mock_core_client_constructor .return_value .close .assert_called_once ()
512+
513+ @patch ("toolbox_langchain.client.ToolboxCoreSyncClient" )
514+ def test_close (self , mock_core_client_constructor ):
515+ """Tests the close method."""
516+ client = ToolboxClient (URL )
517+ client .close ()
518+ mock_core_client_constructor .return_value .close .assert_called_once ()
0 commit comments