chore: Update async client unit tests

Author: anubhav756

packages/toolbox-langchain/tests/test_async_client.py

@@ -18,10 +18,8 @@
 import pytest
 from aiohttp import ClientSession
 from toolbox_core.client import ToolboxClient as ToolboxCoreClient
-from toolbox_core.protocol import ManifestSchema
 from toolbox_core.protocol import ParameterSchema as CoreParameterSchema
 from toolbox_core.tool import ToolboxTool as ToolboxCoreTool
-from toolbox_core.utils import params_to_pydantic_model
 
 from toolbox_langchain.async_client import AsyncToolboxClient
 from toolbox_langchain.async_tools import AsyncToolboxTool
@@ -56,16 +54,12 @@
 
 @pytest.mark.asyncio
 class TestAsyncToolboxClient:
-    @pytest.fixture()
-    def manifest_schema(self):
-        return ManifestSchema(**MANIFEST_JSON)
-
     @pytest.fixture()
     def mock_session(self):
         return AsyncMock(spec=ClientSession)
 
     @pytest.fixture
-    def mock_core_client_instance(self, manifest_schema, mock_session):
+    def mock_core_client_instance(self, mock_session):
         mock = AsyncMock(spec=ToolboxCoreClient)
 
         async def mock_load_tool_impl(name, auth_token_getters, bound_params):
@@ -80,7 +74,8 @@ async def mock_load_tool_impl(name, auth_token_getters, bound_params):
             core_tool_mock = AsyncMock(spec=ToolboxCoreTool)
             core_tool_mock.__name__ = name
             core_tool_mock.__doc__ = tool_schema_dict["description"]
-            core_tool_mock._pydantic_model = params_to_pydantic_model(name, core_params)
+            core_tool_mock._name = name
+            core_tool_mock._params = core_params
             # Add other necessary attributes or method mocks if AsyncToolboxTool uses them
             return core_tool_mock
 
@@ -97,9 +92,8 @@ async def mock_load_toolset_impl(
                 core_tool_mock = AsyncMock(spec=ToolboxCoreTool)
                 core_tool_mock.__name__ = tool_name_iter
                 core_tool_mock.__doc__ = tool_schema_dict["description"]
-                core_tool_mock._pydantic_model = params_to_pydantic_model(
-                    tool_name_iter, core_params
-                )
+                core_tool_mock._name = tool_name_iter
+                core_tool_mock._params = core_params
                 core_tools_list.append(core_tool_mock)
             return core_tools_list
 
@@ -130,86 +124,122 @@ async def test_create_with_existing_session(self, mock_client, mock_session):
     async def test_aload_tool(
         self,
         mock_client,
-        manifest_schema,  # mock_session removed as it's part of mock_core_client_instance
     ):
         tool_name = "test_tool_1"
-        # manifest_schema is used by mock_core_client_instance fixture to provide tool details
+        test_bound_params = {"bp1": "value1"}
 
-        tool = await mock_client.aload_tool(tool_name)
+        tool = await mock_client.aload_tool(tool_name, bound_params=test_bound_params)
 
         # Assert that the core client's load_tool was called correctly
         mock_client._AsyncToolboxClient__core_client.load_tool.assert_called_once_with(
-            name=tool_name, auth_token_getters={}, bound_params={}
+            name=tool_name, auth_token_getters={}, bound_params=test_bound_params
         )
         assert isinstance(tool, AsyncToolboxTool)
         assert (
             tool.name == tool_name
         )  # AsyncToolboxTool gets its name from the core_tool
 
-    async def test_aload_tool_auth_headers_deprecated(
-        self, mock_client, manifest_schema
-    ):
+    async def test_aload_tool_auth_headers_deprecated(self, mock_client):
         tool_name = "test_tool_1"
-        auth_lambda = lambda: "Bearer token"  # Define lambda once
+        auth_lambda = lambda: "Bearer token"
         with catch_warnings(record=True) as w:
             simplefilter("always")
             await mock_client.aload_tool(
                 tool_name,
-                auth_headers={"Authorization": auth_lambda},  # Use the defined lambda
+                auth_headers={"Authorization": auth_lambda},
             )
             assert len(w) == 1
             assert issubclass(w[-1].category, DeprecationWarning)
             assert "auth_headers" in str(w[-1].message)
+            assert "Use `auth_token_getters` instead" in str(w[-1].message)
 
         mock_client._AsyncToolboxClient__core_client.load_tool.assert_called_once_with(
             name=tool_name,
             auth_token_getters={"Authorization": auth_lambda},
             bound_params={},
         )
 
-    async def test_aload_tool_auth_headers_and_tokens(
-        self, mock_client, manifest_schema
-    ):
+    async def test_aload_tool_auth_headers_and_getters_precedence(self, mock_client):
         tool_name = "test_tool_1"
-        auth_getters = {"test": lambda: "token"}
-        auth_headers_lambda = lambda: "Bearer token"  # Define lambda once
+        auth_getters = {"test_source": lambda: "id_token_from_getters"}
+        auth_headers_lambda = lambda: "Bearer token_from_headers"
 
         with catch_warnings(record=True) as w:
             simplefilter("always")
             await mock_client.aload_tool(
                 tool_name,
-                auth_headers={
-                    "Authorization": auth_headers_lambda
-                },  # Use defined lambda
+                auth_headers={"Authorization": auth_headers_lambda},
                 auth_token_getters=auth_getters,
             )
-            assert (
-                len(w) == 1
-            )  # Only one warning because auth_token_getters takes precedence
+            assert len(w) == 1
             assert issubclass(w[-1].category, DeprecationWarning)
-            assert "auth_headers" in str(w[-1].message)  # Warning for auth_headers
+            assert "auth_headers" in str(w[-1].message)
+            assert "`auth_token_getters` will be used" in str(w[-1].message)
 
         mock_client._AsyncToolboxClient__core_client.load_tool.assert_called_once_with(
             name=tool_name, auth_token_getters=auth_getters, bound_params={}
         )
 
-    async def test_aload_toolset(
-        self, mock_client, manifest_schema  # mock_session removed
-    ):
-        tools = await mock_client.aload_toolset()
+    async def test_aload_tool_auth_tokens_deprecated(self, mock_client):
+        tool_name = "test_tool_1"
+        token_lambda = lambda: "id_token"
+        with catch_warnings(record=True) as w:
+            simplefilter("always")
+            await mock_client.aload_tool(
+                tool_name,
+                auth_tokens={"some_token_key": token_lambda},
+            )
+            assert len(w) == 1
+            assert issubclass(w[-1].category, DeprecationWarning)
+            assert "auth_tokens" in str(w[-1].message)
+            assert "Use `auth_token_getters` instead" in str(w[-1].message)
+
+        mock_client._AsyncToolboxClient__core_client.load_tool.assert_called_once_with(
+            name=tool_name,
+            auth_token_getters={"some_token_key": token_lambda},
+            bound_params={},
+        )
+
+    async def test_aload_tool_auth_tokens_and_getters_precedence(self, mock_client):
+        tool_name = "test_tool_1"
+        auth_getters = {"real_source": lambda: "token_from_getters"}
+        token_lambda = lambda: "token_from_auth_tokens"
+
+        with catch_warnings(record=True) as w:
+            simplefilter("always")
+            await mock_client.aload_tool(
+                tool_name,
+                auth_tokens={"deprecated_source": token_lambda},
+                auth_token_getters=auth_getters,
+            )
+            assert len(w) == 1
+            assert issubclass(w[-1].category, DeprecationWarning)
+            assert "auth_tokens" in str(w[-1].message)
+            assert "`auth_token_getters` will be used" in str(w[-1].message)
+
+        mock_client._AsyncToolboxClient__core_client.load_tool.assert_called_once_with(
+            name=tool_name, auth_token_getters=auth_getters, bound_params={}
+        )
+
+    async def test_aload_toolset(self, mock_client):
+        test_bound_params = {"bp_set": "value_set"}
+        tools = await mock_client.aload_toolset(
+            bound_params=test_bound_params, strict=True
+        )
 
         mock_client._AsyncToolboxClient__core_client.load_toolset.assert_called_once_with(
-            name=None, auth_token_getters={}, bound_params={}, strict=False
+            name=None,
+            auth_token_getters={},
+            bound_params=test_bound_params,
+            strict=True,
         )
-        assert len(tools) == 2  # Based on MANIFEST_JSON
+        assert len(tools) == 2
         for tool in tools:
             assert isinstance(tool, AsyncToolboxTool)
             assert tool.name in ["test_tool_1", "test_tool_2"]
 
-    async def test_aload_toolset_with_toolset_name(
-        self, mock_client, manifest_schema  # mock_session removed
-    ):
-        toolset_name = "test_toolset_1"  # This name isn't in MANIFEST_JSON, but load_toolset mock doesn't filter by it
+    async def test_aload_toolset_with_toolset_name(self, mock_client):
+        toolset_name = "test_toolset_1"
         tools = await mock_client.aload_toolset(toolset_name=toolset_name)
 
         mock_client._AsyncToolboxClient__core_client.load_toolset.assert_called_once_with(
@@ -218,43 +248,80 @@ async def test_aload_toolset_with_toolset_name(
         assert len(tools) == 2
         for tool in tools:
             assert isinstance(tool, AsyncToolboxTool)
-            assert tool.name in ["test_tool_1", "test_tool_2"]
 
-    async def test_aload_toolset_auth_headers_deprecated(
-        self, mock_client, manifest_schema
-    ):
-        auth_lambda = lambda: "Bearer token"  # Define lambda once
+    async def test_aload_toolset_auth_headers_deprecated(self, mock_client):
+        auth_lambda = lambda: "Bearer token"
         with catch_warnings(record=True) as w:
             simplefilter("always")
-            await mock_client.aload_toolset(
-                auth_headers={"Authorization": auth_lambda}  # Use defined lambda
-            )
+            await mock_client.aload_toolset(auth_headers={"Authorization": auth_lambda})
             assert len(w) == 1
             assert issubclass(w[-1].category, DeprecationWarning)
             assert "auth_headers" in str(w[-1].message)
+            assert "Use `auth_token_getters` instead" in str(w[-1].message)
+
         mock_client._AsyncToolboxClient__core_client.load_toolset.assert_called_once_with(
             name=None,
             auth_token_getters={"Authorization": auth_lambda},
             bound_params={},
             strict=False,
         )
 
-    async def test_aload_toolset_auth_headers_and_tokens(
-        self, mock_client, manifest_schema
+    async def test_aload_toolset_auth_headers_and_getters_precedence(  # Renamed for clarity
+        self, mock_client
     ):
-        auth_getters = {"test": lambda: "token"}
-        auth_headers_lambda = lambda: "Bearer token"  # Define lambda once
+        auth_getters = {"test_source": lambda: "id_token_from_getters"}
+        auth_headers_lambda = lambda: "Bearer token_from_headers"
         with catch_warnings(record=True) as w:
             simplefilter("always")
             await mock_client.aload_toolset(
-                auth_headers={
-                    "Authorization": auth_headers_lambda
-                },  # Use defined lambda
+                auth_headers={"Authorization": auth_headers_lambda},
                 auth_token_getters=auth_getters,
             )
             assert len(w) == 1
             assert issubclass(w[-1].category, DeprecationWarning)
             assert "auth_headers" in str(w[-1].message)
+            assert "`auth_token_getters` will be used" in str(w[-1].message)
+
+        mock_client._AsyncToolboxClient__core_client.load_toolset.assert_called_once_with(
+            name=None,
+            auth_token_getters=auth_getters,
+            bound_params={},
+            strict=False,  # auth_getters takes precedence
+        )
+
+    async def test_aload_toolset_auth_tokens_deprecated(self, mock_client):
+        token_lambda = lambda: "id_token"
+        with catch_warnings(record=True) as w:
+            simplefilter("always")
+            await mock_client.aload_toolset(
+                auth_tokens={"some_token_key": token_lambda}
+            )
+            assert len(w) == 1
+            assert issubclass(w[-1].category, DeprecationWarning)
+            assert "auth_tokens" in str(w[-1].message)
+            assert "Use `auth_token_getters` instead" in str(w[-1].message)
+
+        mock_client._AsyncToolboxClient__core_client.load_toolset.assert_called_once_with(
+            name=None,
+            auth_token_getters={"some_token_key": token_lambda},
+            bound_params={},
+            strict=False,
+        )
+
+    async def test_aload_toolset_auth_tokens_and_getters_precedence(self, mock_client):
+        auth_getters = {"real_source": lambda: "token_from_getters"}
+        token_lambda = lambda: "token_from_auth_tokens"
+        with catch_warnings(record=True) as w:
+            simplefilter("always")
+            await mock_client.aload_toolset(
+                auth_tokens={"deprecated_source": token_lambda},
+                auth_token_getters=auth_getters,
+            )
+            assert len(w) == 1
+            assert issubclass(w[-1].category, DeprecationWarning)
+            assert "auth_tokens" in str(w[-1].message)
+            assert "`auth_token_getters` will be used" in str(w[-1].message)
+
         mock_client._AsyncToolboxClient__core_client.load_toolset.assert_called_once_with(
             name=None, auth_token_getters=auth_getters, bound_params={}, strict=False
         )