test: upgrade integration harnesses for dual-server execution

packages/toolbox-adk/integration.cloudbuild.yaml

@@ -50,5 +50,5 @@ options:
 substitutions:
   _VERSION: '3.13'
   # Default values (can be overridden by triggers)
-  _TOOLBOX_VERSION: '1.5.0'
+  _TOOLBOX_VERSION: '1.6.0'
   _TOOLBOX_MANIFEST_VERSION: '34'

packages/toolbox-adk/pyproject.toml

@@ -40,7 +40,8 @@ test = [
     "pytest==9.0.3",
     "pytest-asyncio==1.4.0",
     "pytest-cov==7.1.0",
-    "pytest-mock==3.15.1"
+    "pytest-mock==3.15.1",
+    "numpy<2", # Pinned to <2 to prevent mypy syntax errors on python 3.10 with numpy 2.x stubs,
 ]
 
 # Tells setuptools that packages are under the 'src' directory

packages/toolbox-adk/tests/integration/conftest.py

@@ -25,10 +25,13 @@
 from typing import Generator
 
 import google
-import pytest_asyncio
+import pytest
 from google.auth import compute_engine
 from google.cloud import secretmanager, storage
 
+TOOLBOX_SERVER_URL_STABLE = "http://localhost:5000"
+TOOLBOX_SERVER_URL_DRAFT = "http://localhost:5001"
+
 
 #### Define Utility Functions
 def get_env_var(key: str) -> str:
@@ -92,17 +95,17 @@ def get_auth_token(client_id: str) -> str:
 
 
 #### Define Fixtures
-@pytest_asyncio.fixture(scope="session")
+@pytest.fixture(scope="session")
 def project_id() -> str:
     return get_env_var("GOOGLE_CLOUD_PROJECT")
 
 
-@pytest_asyncio.fixture(scope="session")
+@pytest.fixture(scope="session")
 def toolbox_version() -> str:
     return get_env_var("TOOLBOX_VERSION")
 
 
-@pytest_asyncio.fixture(scope="session")
+@pytest.fixture(scope="session")
 def tools_file_path(project_id: str) -> Generator[str]:
     """Provides a temporary file path containing the tools manifest."""
     if os.environ.get("TEST_MOCK_GCP"):
@@ -122,7 +125,7 @@ def tools_file_path(project_id: str) -> Generator[str]:
     os.remove(tools_file_path)
 
 
-@pytest_asyncio.fixture(scope="session")
+@pytest.fixture(scope="session")
 def auth_token1(project_id: str) -> str:
     if os.environ.get("TEST_MOCK_GCP"):
         return "mock-token-1"
@@ -132,7 +135,7 @@ def auth_token1(project_id: str) -> str:
     return get_auth_token(client_id)
 
 
-@pytest_asyncio.fixture(scope="session")
+@pytest.fixture(scope="session")
 def auth_token2(project_id: str) -> str:
     if os.environ.get("TEST_MOCK_GCP"):
         return "mock-token-2"
@@ -142,7 +145,7 @@ def auth_token2(project_id: str) -> str:
     return get_auth_token(client_id)
 
 
-@pytest_asyncio.fixture(scope="session")
+@pytest.fixture(scope="session")
 def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None]:
     """Starts the toolbox server as a subprocess."""
     if os.environ.get("TEST_MOCK_GCP"):
@@ -159,26 +162,62 @@ def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None
         # Make toolbox executable
         os.chmod("toolbox", 0o700)
         # Run toolbox binary
-        toolbox_server = subprocess.Popen(
-            ["./toolbox", "--tools-file", tools_file_path]
+        toolbox_server_1 = subprocess.Popen(
+            ["./toolbox", "--port", "5000", "--tools-file", tools_file_path]
+        )
+        toolbox_server_2 = subprocess.Popen(
+            [
+                "./toolbox",
+                "--port",
+                "5001",
+                "--tools-file",
+                tools_file_path,
+                "--enable-draft-specs",
+            ]
         )
 
         # Wait for server to start
         # Retry logic with a timeout
         for _ in range(5):  # retries
             time.sleep(2)
-            print("Checking if toolbox is successfully started...")
-            if toolbox_server.poll() is None:
-                print("Toolbox server started successfully.")
+            print("Checking if both toolbox servers are successfully started...")
+            if toolbox_server_1.poll() is None and toolbox_server_2.poll() is None:
+                print("Toolbox servers started successfully.")
                 break
         else:
-            raise RuntimeError("Toolbox server failed to start after 5 retries.")
+            raise RuntimeError("Toolbox servers failed to start after 5 retries.")
     except subprocess.CalledProcessError as e:
         print(e.stderr.decode("utf-8"))
         print(e.stdout.decode("utf-8"))
         raise RuntimeError(f"{e}\n\n{e.stderr.decode('utf-8')}") from e
     yield
 
     # Clean up toolbox server
-    toolbox_server.terminate()
-    toolbox_server.wait(timeout=5)
+    toolbox_server_1.terminate()
+    toolbox_server_2.terminate()
+    toolbox_server_1.wait(timeout=5)
+    toolbox_server_2.wait(timeout=5)
+
+
+@pytest.fixture(
+    params=[TOOLBOX_SERVER_URL_STABLE, TOOLBOX_SERVER_URL_DRAFT], scope="session"
+)
+def toolbox_server_url(request) -> str:
+    return request.param
+
+
+@pytest.fixture()
+def patch_toolbox_client_url(toolbox_server_url):
+    from toolbox_adk.toolset import ToolboxToolset
+
+    original_init = ToolboxToolset.__init__
+
+    def new_init(self, server_url=TOOLBOX_SERVER_URL_STABLE, *args, **kwargs):
+        if server_url == TOOLBOX_SERVER_URL_STABLE:
+            server_url = toolbox_server_url
+        original_init(self, server_url, *args, **kwargs)
+
+    from unittest.mock import patch
+
+    with patch.object(ToolboxToolset, "__init__", new_init, create=True):
+        yield