mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-08-18 16:02:10 +00:00
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>
This commit is contained in:
Alexander Petric
GitHub
@@ -5,3 +5,4 @@ httpcore==1.0.2
|
||||
httpx==0.26.0
|
||||
idna==3.6
|
||||
sniffio==1.3.0
|
||||
docker==7.1.0
|
||||
|
||||
@@ -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 *
|
||||
|
||||
@@ -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_<prefix>_<token>
|
||||
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)
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user