From dda92b5491beb41d8a561be3994f65de58f0665e Mon Sep 17 00:00:00 2001 From: Alexander Petric Date: Tue, 29 Apr 2025 17:55:38 -0400 Subject: [PATCH] agent integration tests (#5684) * agent integration tests * agent integration tests * Update integration_tests/requirements.txt Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com> --------- Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com> --- integration_tests/requirements.txt | 1 + integration_tests/test/__init__.py | 1 + integration_tests/test/agent_workers.py | 265 ++++++++++++++++++ .../test/wmill_integration_test_utils.py | 181 +++++++++++- 4 files changed, 440 insertions(+), 8 deletions(-) create mode 100644 integration_tests/test/agent_workers.py diff --git a/integration_tests/requirements.txt b/integration_tests/requirements.txt index 4b36ac91dd..97572f1b84 100644 --- a/integration_tests/requirements.txt +++ b/integration_tests/requirements.txt @@ -5,3 +5,4 @@ httpcore==1.0.2 httpx==0.26.0 idna==3.6 sniffio==1.3.0 +docker==7.1.0 diff --git a/integration_tests/test/__init__.py b/integration_tests/test/__init__.py index adab364c51..6a7af10321 100644 --- a/integration_tests/test/__init__.py +++ b/integration_tests/test/__init__.py @@ -2,3 +2,4 @@ from .identity_script_test import * from .increment_flow_test import * from .schedule_test import * from .windmill_sdk_test import * +from .agent_workers import * diff --git a/integration_tests/test/agent_workers.py b/integration_tests/test/agent_workers.py new file mode 100644 index 0000000000..82a2cf3b52 --- /dev/null +++ b/integration_tests/test/agent_workers.py @@ -0,0 +1,265 @@ +import unittest +import base64 +import json +import os +import time +import atexit +import docker +from docker.errors import NotFound + +from .wmill_integration_test_utils import WindmillClient + +AGENT_CONTAINER_NAME = "windmill_agent_test" + +# Global cleanup function registered with atexit +def cleanup_docker_container(): + print(f"At exit: Attempting to remove container {AGENT_CONTAINER_NAME}") + try: + docker_client = docker.from_env() + try: + container = docker_client.containers.get(AGENT_CONTAINER_NAME) + container.stop() + container.remove() + print(f"At exit: Successfully removed container {AGENT_CONTAINER_NAME}") + except NotFound: + print(f"At exit: Container {AGENT_CONTAINER_NAME} not found") + except Exception as e: + print(f"At exit: Error removing container: {e}") + except Exception as e: + print(f"At exit: Could not initialize Docker client: {e}") + +atexit.register(cleanup_docker_container) + + +class TestAgentWorkers(unittest.TestCase): + _client: WindmillClient + _docker_client = None + _agent_container = None + _container_name = AGENT_CONTAINER_NAME + _agent_token = None + + @classmethod + def setUpClass(cls) -> None: + print("Running {}".format(cls.__name__)) + cls._client = WindmillClient() + + cls._client.add_global_custom_tag("agent_test") + cls._docker_client = docker.from_env() + cls._start_agent_container() + cls._wait_for_agent_connection() + + @classmethod + def _start_agent_container(cls): + wm_image = os.environ.get("WM_IMAGE", "ghcr.io/windmill-labs/windmill-ee") + wm_version = os.environ.get("WM_VERSION", "latest") + + cls._agent_token = cls._client.create_agent_token( + worker_group="agent", + tags=["agent", "python3", "bash", "agent_test"], + exp=int(time.time()) + 3600 # 1 hour from now + ) + + try: + container = cls._docker_client.containers.get(cls._container_name) + container.stop() + container.remove() + print(f"Removed existing container {cls._container_name}") + except NotFound: + pass + + # Get the host IP address for connecting back to the Windmill server + host_ip = "host.docker.internal" + if os.name == "posix" and os.uname().sysname == "Linux": + import socket + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + s.connect(("8.8.8.8", 80)) + host_ip = s.getsockname()[0] + s.close() + + print(f"Starting agent container connecting to Windmill server at http://{host_ip}:8000") + + # Create and start the container + container = cls._docker_client.containers.run( + f"{wm_image}:{wm_version}", + name=cls._container_name, + detach=True, + environment={ + "BASE_INTERNAL_URL": f"http://{host_ip}:8000", + "MODE": "agent", + "AGENT_TOKEN": cls._agent_token + }, + volumes={ + "/var/run/docker.sock": {"bind": "/var/run/docker.sock", "mode": "rw"} + }, + restart_policy={"Name": "unless-stopped"}, + mem_limit="2g", + cpu_count=1, + ) + + cls._agent_container = container + print(f"Started agent container: {cls._container_name}") + + @classmethod + def _wait_for_agent_connection(cls): + print("Waiting for agent to connect to the server...") + connected = False + max_attempts = 60 + for attempt in range(max_attempts): + # Query server for connected workers + workers = cls._client.get_workers_list(ping_since=60) + + # Check if any worker is in the "agent" worker_group + for worker in workers: + if worker.get("worker_group") == "agent": + connected = True + print(f"Agent connected with details: {worker}") + break + + if connected: + break + + print(f"Waiting for agent to connect... Attempt {attempt+1}/{max_attempts}") + time.sleep(1) + + if not connected: + raise Exception("Agent failed to connect within the expected time") + + print("Agent successfully connected!") + + @classmethod + def tearDownClass(cls) -> None: + print("Cleaning up after tests...") + + if hasattr(cls, "_client") and cls._client is not None: + cls._client.remove_global_custom_tag("agent_test") + + if cls._docker_client is None: + try: + cls._docker_client = docker.from_env() + except Exception as e: + print(f"Error initializing Docker client for cleanup: {e}") + return + + if cls._agent_container is not None: + try: + cls._agent_container.stop() + cls._agent_container.remove() + print(f"Cleaned up container using container object: {cls._container_name}") + return + except Exception as e: + print(f"Error cleaning up container using object: {e}") + + try: + container = cls._docker_client.containers.get(cls._container_name) + container.stop() + container.remove() + print(f"Cleaned up container by name: {cls._container_name}") + except NotFound: + print(f"Container {cls._container_name} not found during cleanup") + except Exception as e: + print(f"Error cleaning up container by name: {e}") + + print("Cleanup complete") + + def test_create_agent_token(self): + token = self._agent_token + self.assertIsNotNone(token) + + # JWT tokens have the format: jwt_agent__ + self.assertTrue(token.startswith("jwt_agent_"), "Token should start with jwt_agent_") + + # Test that it's a valid JWT format (should contain 2 dots in the JWT part) + parts = token.split('_') + self.assertGreaterEqual(len(parts), 3, "Token should have at least 3 parts separated by underscores") + + # The actual JWT is after the second underscore + jwt_part = parts[2] + self.assertEqual(jwt_part.count('.'), 2, "JWT should contain exactly 2 dots") + + # Check that the token contains three base64-encoded parts + jwt_segments = jwt_part.split('.') + self.assertEqual(len(jwt_segments), 3, "JWT should have 3 segments") + for segment in jwt_segments: + self.assertGreater(len(segment), 0, "JWT segment should not be empty") + + # Decode the JWT payload (second segment) + payload_base64 = jwt_segments[1] + + # Add padding if needed + padding_needed = len(payload_base64) % 4 + if padding_needed: + payload_base64 += '=' * (4 - padding_needed) + + # Base64 decode the payload + payload_bytes = base64.urlsafe_b64decode(payload_base64) + payload_json = payload_bytes.decode('utf-8') + payload = json.loads(payload_json) + + # Check payload structure + self.assertIn('worker_group', payload, "Payload should contain worker_group") + self.assertEqual(payload['worker_group'], 'agent', "worker_group should be 'agent'") + + self.assertIn('suffix', payload, "Payload should contain suffix") + self.assertIsNone(payload['suffix'], "suffix should be null") + + self.assertIn('tags', payload, "Payload should contain tags") + self.assertIsInstance(payload['tags'], list, "tags should be a list") + self.assertIn('agent', payload['tags'], "tags should contain 'agent'") + + self.assertIn('exp', payload, "Payload should contain exp") + self.assertIsInstance(payload['exp'], int, "exp should be an integer") + + def test_agent_is_connected(self): + """Test that the agent is connected to the server.""" + workers = self._client.get_workers_list(ping_since=60) + + # Find the agent worker + agent_worker = None + for worker in workers: + if worker.get("worker_group") == "agent": + agent_worker = worker + break + + self.assertIsNotNone(agent_worker, "Agent worker should be connected") + self.assertEqual(agent_worker.get("worker_group"), "agent") + + # Check tags + tags = agent_worker.get("custom_tags", []) + self.assertIn("agent", tags) + self.assertIn("python3", tags) + self.assertIn("bash", tags) + self.assertIn("agent_test", tags) + + def test_bash_script_on_agent(self): + """Test running a bash script on the agent worker.""" + # Create a simple bash script tagged to run on the agent + script_path = "u/admin/agent_bash_test" + script_content = """ +#!/bin/bash +msg="$1" +echo "Running on $(hostname)" +echo "Argument received: $msg" +echo $msg +""" + + # Create the script with the agent_test tag so it runs on our agent + self._client.create_script( + path=script_path, + content=script_content, + language="bash", + tag="agent_test" + ) + + try: + # Run the script + result = self._client.run_sync(script_path, {"msg": "Hello from agent test!"}) + + print(f"Script result: {result}") + + # Verify the result + self.assertIsNotNone(result) + self.assertEqual(result, "Hello from agent test!") + + finally: + # Clean up the script + self._client.delete_script(script_path) \ No newline at end of file diff --git a/integration_tests/test/wmill_integration_test_utils.py b/integration_tests/test/wmill_integration_test_utils.py index cd49e7c1ec..bb4b2ffe22 100644 --- a/integration_tests/test/wmill_integration_test_utils.py +++ b/integration_tests/test/wmill_integration_test_utils.py @@ -114,17 +114,23 @@ class WindmillClient: raise Exception(response.content.decode()) return response.json() - def create_script(self, path: str, content: str, language: str): + def create_script(self, path: str, content: str, language: str, tag: str = None): print(f"Creating script {path}") + + payload = { + "path": path, + "content": content, + "description": "", + "summary": "", + "language": language, + } + + if tag is not None: + payload["tag"] = tag + response = self._client.post( f"/api/w/{self._workspace}/scripts/create", - json={ - "path": path, - "content": content, - "description": "", - "summary": "", - "language": language, - }, + json=payload, ) if response.status_code // 100 != 2: raise Exception(response.content.decode()) @@ -252,3 +258,162 @@ class WindmillClient: def get_version(self): response = self._client.get("/api/version") return response.content.decode() + + def get_global_custom_tags(self): + """ + Get the current list of global custom tags. + + Returns: + list: List of custom tags or empty list if not set or an error occurred. + """ + try: + response = self._client.get("/api/settings/global/custom_tags") + if response.status_code // 100 == 2: + tags = response.json() + return tags if tags is not None else [] + else: + print(f"Error retrieving global custom tags: Status {response.status_code}, Response: {response.content.decode()}") + return [] + except Exception as e: + print(f"Exception when retrieving global custom tags: {e}") + return [] + + def add_global_custom_tag(self, tag): + """ + Add a tag to the global custom tags if it's not already present. + + Args: + tag (str): The tag to add to global custom tags. + + Returns: + bool: True if the tag was added or already exists, False if there was an error. + """ + try: + current_tags = self.get_global_custom_tags() + + if tag in current_tags: + print(f"Tag '{tag}' already exists in global custom tags") + return True + + new_tags = current_tags + [tag] + print(f"Adding '{tag}' to global custom tags: {new_tags}") + + response = self._client.post( + "/api/settings/global/custom_tags", + json={ + "value": new_tags, + }, + ) + + if response.status_code // 100 == 2: + print(f"Successfully added '{tag}' to global custom tags") + return True + else: + print(f"Error adding tag to global custom tags: Status {response.status_code}, Response: {response.content.decode()}") + return False + except Exception as e: + print(f"Exception when adding global custom tag: {e}") + return False + + def remove_global_custom_tag(self, tag): + """ + Remove a tag from the global custom tags if it exists. + + Args: + tag (str): The tag to remove from global custom tags. + + Returns: + bool: True if the tag was removed or didn't exist, False if there was an error. + """ + try: + current_tags = self.get_global_custom_tags() + + if tag not in current_tags: + print(f"Tag '{tag}' doesn't exist in global custom tags") + return True + + new_tags = [t for t in current_tags if t != tag] + print(f"Removing '{tag}' from global custom tags: {new_tags}") + + response = self._client.post( + "/api/settings/global/custom_tags", + json={ + "value": new_tags, + }, + ) + + if response.status_code // 100 == 2: + print(f"Successfully removed '{tag}' from global custom tags") + return True + else: + print(f"Error removing tag from global custom tags: Status {response.status_code}, Response: {response.content.decode()}") + return False + except Exception as e: + print(f"Exception when removing global custom tag: {e}") + return False + + def get_workers_list(self, ping_since=60, page=0, per_page=100): + """ + Get a list of workers currently connected to the Windmill server. + + Args: + ping_since (int): Only include workers that have pinged in the last N seconds. Default is 60. + page (int): Page number for pagination. Default is 0. + per_page (int): Number of results per page. Default is 100. + + Returns: + list: List of worker objects or empty list if no workers found or an error occurred. + """ + try: + params = { + "page": page, + "per_page": per_page, + "ping_since": ping_since + } + + response = self._client.get( + "/api/workers/list", + params=params + ) + if response.status_code // 100 == 2: + return response.json() + else: + print(f"Error retrieving workers list: Status {response.status_code}, Response: {response.content.decode()}") + return [] + except Exception as e: + print(f"Exception when retrieving workers list: {e}") + return [] + + def create_agent_token(self, worker_group="agent", tags=None, exp=None): + """ + Create an agent JWT token using superadmin privilege. + + Args: + worker_group (str): The worker group for the agent, defaults to "agent" + tags (list): Tags for the agent, defaults to ["agent"] + exp (int): Expiration timestamp, defaults to a timestamp about 1 year in the future + + Returns: + str: The JWT token for the agent + """ + if tags is None: + tags = ["agent"] + + if exp is None: + exp = int(time.time()) + 31536000 # 60*60*24*365 = 1 year + + print(f"Creating agent token for worker_group={worker_group}, tags={tags}") + response = self._client.post( + "/api/agent_workers/create_agent_token", + json={ + "worker_group": worker_group, + "tags": tags, + "exp": exp + }, + ) + if response.status_code // 100 != 2: + raise Exception(response.content.decode()) + + token = response.content.decode().strip('"') + print(f"Created agent token: {token[:15]}...{token[-15:]}") + return token