import unittest
import subprocess
from subprocess import TimeoutExpired
import sys
import json
import os
from pathlib import Path
import pytest
class TestMCPIntegration(unittest.TestCase):
def setUp(self):
self.repo_root = Path(__file__).resolve().parent.parent
env = os.environ.copy()
env["PYTHONPATH"] = str(self.repo_root / "src") + os.pathsep + env.get("PYTHONPATH", "")
self.env = env
# Use python -m to run MCP server directly without CLI dependencies
self.base_cmd = [
sys.executable,
"-m",
"zimage.mcp_server",
]
def run_process(self, input_str):
process = subprocess.Popen(
self.base_cmd,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
env=self.env,
cwd=self.repo_root,
)
try:
stdout, stderr = process.communicate(input=input_str, timeout=20)
except TimeoutExpired:
# If the MCP server stays running after handling the request, terminate to avoid hanging tests
process.terminate()
try:
stdout, stderr = process.communicate(timeout=5)
except TimeoutExpired:
process.kill()
stdout, stderr = process.communicate(timeout=5)
return stdout, stderr
def test_initialize(self):
req = {
"jsonrpc": "2.0",
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
},
"id": 1
}
stdout, stderr = self.run_process(json.dumps(req) + "\n")
response = None
for line in stdout.splitlines():
try:
msg = json.loads(line)
if msg.get("id") == 1:
response = msg
break
except json.JSONDecodeError:
continue
self.assertIsNotNone(response, f"Did not receive JSON-RPC response. Stdout: {stdout}\nStderr: {stderr}")
self.assertIn("result", response)
self.assertIn("serverInfo", response["result"])
self.assertEqual(response["result"]["serverInfo"]["name"], "Z-Image Studio")
def test_call_list_models(self):
# Send initialize -> initialized -> tools/call
req_init = {
"jsonrpc": "2.0",
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
},
"id": 1
}
req_notify = {
"jsonrpc": "2.0",
"method": "notifications/initialized",
"params": {}
}
req_call = {
"jsonrpc": "2.0",
"method": "tools/call",
"params": {
"name": "list_models",
"arguments": {}
},
"id": 2
}
input_str = json.dumps(req_init) + "\n" + json.dumps(req_notify) + "\n" + json.dumps(req_call) + "\n"
stdout, stderr = self.run_process(input_str)
responses = []
for line in stdout.splitlines():
try:
msg = json.loads(line)
responses.append(msg)
except:
pass
# Find response with id 2
res = next((r for r in responses if r.get("id") == 2), None)
self.assertIsNotNone(res, f"Did not receive call_tool response. Stdout: {stdout}\nStderr: {stderr}")
if "error" in res:
self.fail(f"Received error from tool: {res['error']}")
self.assertIn("result", res)
# Check if output contains text
content = res["result"].get("content", [])
self.assertTrue(len(content) > 0)
text = content[0].get("text", "")
self.assertIn("Available Models", text)
def test_seed_generation_logic(self):
"""Test the seed generation logic without requiring the full model."""
# This test verifies the core seed generation logic that was added
import random
# Simulate the seed generation logic from mcp_server.py
def simulate_seed_generation(seed_input):
if seed_input is None:
return random.randint(0, 2**31 - 1)
else:
return seed_input
# Test case 1: None input should generate a random seed
seed1 = simulate_seed_generation(None)
self.assertIsNotNone(seed1, "Seed should not be None")
self.assertIsInstance(seed1, int, "Seed should be an integer")
self.assertGreaterEqual(seed1, 0, "Seed should be non-negative")
self.assertLessEqual(seed1, 2**31 - 1, "Seed should be within valid range")
# Test case 2: Provided seed should be used as-is
provided_seed = 12345
seed2 = simulate_seed_generation(provided_seed)
self.assertEqual(seed2, provided_seed, "Should use provided seed")
# Test case 3: Zero seed should be preserved
zero_seed = 0
seed3 = simulate_seed_generation(zero_seed)
self.assertEqual(seed3, zero_seed, "Should preserve zero seed")
def test_call_generate_with_null_seed_lightweight(self):
"""Test seed generation logic in isolation (not real endpoint)."""
# This test verifies that the seed generation logic works correctly
# by testing the actual code path that would be executed
import random
# Simulate the exact seed generation logic from mcp_server.py generate function
def simulate_mcp_generate_with_seed_handling(seed_input):
"""Simulate the seed handling logic from the actual MCP server."""
# This is the exact logic from mcp_server.py
if seed_input is None:
seed = random.randint(0, 2**31 - 1)
return seed
else:
return seed_input
# Test the seed generation with None input
seed_input = None
generated_seed = simulate_mcp_generate_with_seed_handling(seed_input)
# Verify the generated seed meets all requirements
self.assertIsNotNone(generated_seed, "Generated seed should not be None")
self.assertIsInstance(generated_seed, int, "Generated seed should be an integer")
self.assertGreaterEqual(generated_seed, 0, "Generated seed should be non-negative")
self.assertLessEqual(generated_seed, 2**31 - 1, f"Generated seed {generated_seed} should be within valid range")
# Test with provided seed
provided_seed = 12345
result_seed = simulate_mcp_generate_with_seed_handling(provided_seed)
self.assertEqual(result_seed, provided_seed, "Should use provided seed when not None")
# Test edge case: seed = 0
zero_seed = 0
result_seed = simulate_mcp_generate_with_seed_handling(zero_seed)
self.assertEqual(result_seed, zero_seed, "Should preserve zero seed")
@unittest.skip("This test requires the full model and GPU, skipping to avoid flaky CI")
def test_call_generate_with_null_seed(self):
"""Test that generate tool creates a seed when none is provided."""
# This test is skipped by default as it requires the full model and GPU
# It's kept here for documentation and manual testing purposes
# Send initialize -> initialized -> tools/call with null seed
req_init = {
"jsonrpc": "2.0",
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
},
"id": 1
}
req_notify = {
"jsonrpc": "2.0",
"method": "notifications/initialized",
"params": {}
}
req_call = {
"jsonrpc": "2.0",
"method": "tools/call",
"params": {
"name": "generate",
"arguments": {
"prompt": "a test image",
"steps": 2, # Minimal steps for faster test
"width": 256,
"height": 256,
"seed": None # Explicitly null seed
}
},
"id": 2
}
input_str = json.dumps(req_init) + "\n" + json.dumps(req_notify) + "\n" + json.dumps(req_call) + "\n"
stdout, stderr = self.run_process(input_str)
responses = []
for line in stdout.splitlines():
try:
msg = json.loads(line)
responses.append(msg)
except:
pass
# Find response with id 2
res = next((r for r in responses if r.get("id") == 2), None)
self.assertIsNotNone(res, f"Did not receive call_tool response. Stdout: {stdout}\nStderr: {stderr}")
if "error" in res:
self.fail(f"Received error from tool: {res['error']}")
self.assertIn("result", res)
# Check if output contains content
content = res["result"].get("content", [])
self.assertTrue(len(content) > 0)
# Find the text content (metadata)
text_content = next((c for c in content if c.get("type") == "text"), None)
self.assertIsNotNone(text_content, "No text content found in response")
# Parse the metadata
metadata = json.loads(text_content["text"])
# The key assertion: seed should NOT be null, it should be a generated integer
self.assertIn("seed", metadata, "Seed not found in metadata")
self.assertIsNotNone(metadata["seed"], "Seed should not be null")
self.assertIsInstance(metadata["seed"], int, f"Seed should be an integer, got {type(metadata['seed'])}")
self.assertGreaterEqual(metadata["seed"], 0, "Seed should be non-negative")
self.assertLessEqual(metadata["seed"], 2**31 - 1, f"Seed {metadata['seed']} should be within valid range")
@unittest.skip("This test requires the full model and GPU, skipping to avoid flaky CI")
def test_transport_content_consistency(self):
"""Test that stdio and SSE transports return identical content structure."""
# This test verifies the fix for issue #35 - transport-agnostic content
# Test stdio transport
req_init = {
"jsonrpc": "2.0",
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {"name": "test", "version": "1.0"}
},
"id": 1
}
req_notify = {
"jsonrpc": "2.0",
"method": "notifications/initialized",
"params": {}
}
req_call = {
"jsonrpc": "2.0",
"method": "tools/call",
"params": {
"name": "generate",
"arguments": {
"prompt": "a test image",
"steps": 2,
"width": 256,
"height": 256,
"seed": 12345
}
},
"id": 2
}
input_str = json.dumps(req_init) + "\n" + json.dumps(req_notify) + "\n" + json.dumps(req_call) + "\n"
stdout, stderr = self.run_process(input_str)
responses = []
for line in stdout.splitlines():
try:
msg = json.loads(line)
responses.append(msg)
except:
pass
# Find response with id 2
res = next((r for r in responses if r.get("id") == 2), None)
self.assertIsNotNone(res, f"Did not receive call_tool response. Stdout: {stdout}\nStderr: {stderr}")
if "error" in res:
self.fail(f"Received error from tool: {res['error']}")
self.assertIn("result", res)
content = res["result"].get("content", [])
# Verify content structure: should have exactly 3 items in order:
# 1. TextContent (metadata)
# 2. ResourceContent (main image file)
# 3. ImageContent (thumbnail)
self.assertEqual(len(content), 3, f"Expected 3 content items, got {len(content)}")
# Check content types
self.assertEqual(content[0]["type"], "text", "First content should be text")
self.assertEqual(content[1]["type"], "resource", "Second content should be resource")
self.assertEqual(content[2]["type"], "image", "Third content should be image")
# Verify text content structure (now includes file metadata)
text_content = content[0]
self.assertIn("text", text_content)
metadata = json.loads(text_content["text"])
self.assertIn("seed", metadata)
self.assertIn("duration_seconds", metadata)
self.assertIn("width", metadata)
self.assertIn("height", metadata)
self.assertIn("filename", metadata)
self.assertIn("file_path", metadata)
# Verify resource link structure (clean URI only, no _meta)
resource_content = content[1]
self.assertEqual(resource_content["type"], "resource_link")
self.assertIn("uri", resource_content)
self.assertIn("mimeType", resource_content)
self.assertIn("name", resource_content)
self.assertNotIn("_meta", resource_content, "ResourceLink should not have _meta")
# Verify stdio transport uses file:// URI
self.assertTrue(resource_content["uri"].startswith("file://"), f"Stdio should use file:// URI, got: {resource_content['uri']}")
# Verify image content structure
image_content = content[2]
self.assertIn("data", image_content)
self.assertIn("mimeType", image_content)
self.assertEqual(image_content["mimeType"], "image/png")
# Verify file metadata is properly included in text content (not duplicated)
text_metadata = json.loads(content[0]["text"])
self.assertIn("file_path", text_metadata, "file_path should be in text content")
self.assertIn("filename", text_metadata, "filename should be in text content")
# Ensure no URL duplication (URLs should only be in ResourceContent)
self.assertNotIn("relative_url", text_metadata, "relative_url should not be in text content")
self.assertNotIn("absolute_url", text_metadata, "absolute_url should not be in text content")
if __name__ == '__main__':
unittest.main()
