"""Tests for the Yellhorn MCP server."""
import asyncio
import json
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from google import genai
from pydantic import FileUrl
@pytest.mark.asyncio
async def test_list_resources(mock_request_context):
"""Test listing workplan and judgement sub-issue resources."""
with (
patch("yellhorn_mcp.utils.git_utils.run_github_command") as mock_gh,
patch("yellhorn_mcp.utils.git_utils.Resource") as mock_resource_class,
):
# Set up 2 workplan issues and 2 review sub-issues
# Configure mock responses for different labels
def mock_gh_side_effect(*args, **kwargs):
if "--label" in args[1] and args[1][args[1].index("--label") + 1] == "yellhorn-mcp":
return """[
{"number": 123, "title": "Test Workplan 1", "url": "https://github.com/user/repo/issues/123"},
{"number": 456, "title": "Test Workplan 2", "url": "https://github.com/user/repo/issues/456"}
]"""
elif (
"--label" in args[1]
and args[1][args[1].index("--label") + 1] == "yellhorn-judgement-subissue"
):
return """[
{"number": 789, "title": "Judgement: main..HEAD for Workplan #123", "url": "https://github.com/user/repo/issues/789"},
{"number": 987, "title": "Judgement: v1.0..feature for Workplan #456", "url": "https://github.com/user/repo/issues/987"}
]"""
return "[]"
mock_gh.side_effect = mock_gh_side_effect
# Configure mock_resource_class to return mock Resource objects
workplan1 = MagicMock()
workplan1.uri = FileUrl(f"file://workplans/123.md")
workplan1.name = "Workplan #123: Test Workplan 1"
workplan1.mimeType = "text/markdown"
workplan2 = MagicMock()
workplan2.uri = FileUrl(f"file://workplans/456.md")
workplan2.name = "Workplan #456: Test Workplan 2"
workplan2.mimeType = "text/markdown"
judgement1 = MagicMock()
judgement1.uri = FileUrl(f"file://judgements/789.md")
judgement1.name = "Judgement #789: Judgement: main..HEAD for Workplan #123"
judgement1.mimeType = "text/markdown"
judgement2 = MagicMock()
judgement2.uri = FileUrl(f"file://judgements/987.md")
judgement2.name = "Judgement #987: Judgement: v1.0..feature for Workplan #456"
judgement2.mimeType = "text/markdown"
# Configure the Resource constructor to return our mock objects
mock_resource_class.side_effect = [workplan1, workplan2, judgement1, judgement2]
# 1. Test with no resource_type (should get both types)
resources = await list_resources(mock_request_context, None)
# Verify the GitHub command was called correctly for both labels
assert mock_gh.call_count == 2
mock_gh.assert_any_call(
mock_request_context.request_context.lifespan_context["repo_path"],
["issue", "list", "--label", "yellhorn-mcp", "--json", "number,title,url"],
)
mock_gh.assert_any_call(
mock_request_context.request_context.lifespan_context["repo_path"],
[
"issue",
"list",
"--label",
"yellhorn-judgement-subissue",
"--json",
"number,title,url",
],
)
# Verify Resource constructor was called for all 4 resources
assert mock_resource_class.call_count == 4
# Verify resources are returned correctly (both types)
assert len(resources) == 4
# Reset mocks for the next test
mock_gh.reset_mock()
mock_resource_class.reset_mock()
mock_resource_class.side_effect = [workplan1, workplan2]
# 2. Test with resource_type="yellhorn_workplan" - should return only workplans
resources = await list_resources(mock_request_context, "yellhorn_workplan")
assert len(resources) == 2
mock_gh.assert_called_once_with(
mock_request_context.request_context.lifespan_context["repo_path"],
["issue", "list", "--label", "yellhorn-mcp", "--json", "number,title,url"],
)
# Reset mocks for the next test
mock_gh.reset_mock()
mock_resource_class.reset_mock()
mock_resource_class.side_effect = [judgement1, judgement2]
# 3. Test with resource_type="yellhorn_judgement_subissue" - should return only judgements
resources = await list_resources(mock_request_context, "yellhorn_judgement_subissue")
assert len(resources) == 2
mock_gh.assert_called_once_with(
mock_request_context.request_context.lifespan_context["repo_path"],
[
"issue",
"list",
"--label",
"yellhorn-judgement-subissue",
"--json",
"number,title,url",
],
)
# Reset mock for the final test
mock_gh.reset_mock()
# 4. Test with different resource_type - should return empty list
resources = await list_resources(mock_request_context, "different_type")
assert len(resources) == 0
# GitHub command should not be called in this case
mock_gh.assert_not_called()
@pytest.mark.asyncio
async def test_read_resource(mock_request_context):
"""Test getting resources by type."""
with patch("yellhorn_mcp.utils.git_utils.get_github_issue_body") as mock_get_issue:
# Test 1: Get workplan resource
mock_get_issue.return_value = "# Test Workplan\n\n1. Step one\n2. Step two"
# Call the read_resource method with yellhorn_workplan type
resource_content = await read_resource(mock_request_context, "123", "yellhorn_workplan")
# Verify the GitHub issue body was retrieved correctly
mock_get_issue.assert_called_once_with(
mock_request_context.request_context.lifespan_context["repo_path"], "123"
)
# Verify resource content is returned correctly
assert resource_content == "# Test Workplan\n\n1. Step one\n2. Step two"
# Reset mock for next test
mock_get_issue.reset_mock()
# Test 2: Get judgement sub-issue resource
mock_get_issue.return_value = (
"## Judgement Summary\nThis is a judgement of the implementation."
)
# Call the read_resource method with yellhorn_judgement_subissue type
resource_content = await read_resource(
mock_request_context, "456", "yellhorn_judgement_subissue"
)
# Verify the GitHub issue body was retrieved correctly
mock_get_issue.assert_called_once_with(
mock_request_context.request_context.lifespan_context["repo_path"], "456"
)
# Verify resource content is returned correctly
assert (
resource_content == "## Judgement Summary\nThis is a judgement of the implementation."
)
# Reset mock for next test
mock_get_issue.reset_mock()
# Test 3: Get resource without specifying type
mock_get_issue.return_value = "# Any content"
# Call the read_resource method without type
resource_content = await read_resource(mock_request_context, "789", None)
# Verify the GitHub issue body was retrieved correctly
mock_get_issue.assert_called_once_with(
mock_request_context.request_context.lifespan_context["repo_path"], "789"
)
# Verify resource content is returned correctly
assert resource_content == "# Any content"
# Test with unsupported resource type
with pytest.raises(ValueError, match="Unsupported resource type"):
await read_resource(mock_request_context, "123", "unsupported_type")
from mcp.server.fastmcp import Context
from yellhorn_mcp.server import (
calculate_cost,
create_workplan,
format_metrics_section,
get_codebase_snapshot,
get_git_diff,
get_workplan,
judge_workplan,
process_workplan_async,
revise_workplan,
)
from yellhorn_mcp.utils.git_utils import (
YellhornMCPError,
add_github_issue_comment,
create_github_subissue,
ensure_label_exists,
get_default_branch,
get_github_issue_body,
get_github_pr_diff,
is_git_repository,
list_resources,
post_github_pr_review,
read_resource,
run_git_command,
run_github_command,
update_github_issue,
)
@pytest.fixture
def mock_request_context():
"""Fixture for mock request context."""
mock_ctx = MagicMock(spec=Context)
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/mock/repo"),
"gemini_client": MagicMock(spec=genai.Client),
"openai_client": None,
"model": "gemini-2.5-pro",
}
return mock_ctx
@pytest.fixture
def mock_genai_client():
"""Fixture for mock Gemini API client."""
client = MagicMock(spec=genai.Client)
response = MagicMock()
response.text = "Mock response text"
client.aio.models.generate_content = AsyncMock(return_value=response)
return client
@pytest.mark.asyncio
async def test_run_git_command_success():
"""Test successful Git command execution."""
with patch("asyncio.create_subprocess_exec") as mock_exec:
mock_process = AsyncMock()
mock_process.communicate.return_value = (b"output", b"")
mock_process.returncode = 0
mock_exec.return_value = mock_process
result = await run_git_command(Path("/mock/repo"), ["status"])
assert result == "output"
mock_exec.assert_called_once()
@pytest.mark.asyncio
async def test_run_git_command_failure():
"""Test failed Git command execution."""
with patch("asyncio.create_subprocess_exec") as mock_exec:
mock_process = AsyncMock()
mock_process.communicate.return_value = (b"", b"error message")
mock_process.returncode = 1
mock_exec.return_value = mock_process
with pytest.raises(YellhornMCPError, match="Git command failed: error message"):
await run_git_command(Path("/mock/repo"), ["status"])
@pytest.mark.asyncio
async def test_get_codebase_snapshot():
"""Test getting codebase snapshot."""
with patch("yellhorn_mcp.utils.git_utils.run_git_command") as mock_git:
# Mock both calls to run_git_command (tracked files and untracked files)
mock_git.side_effect = [
"file1.py\nfile2.py", # tracked files
"file3.py", # untracked files
]
with patch("pathlib.Path.is_dir", return_value=False):
with patch("pathlib.Path.exists", return_value=False):
# Mock Path.stat() for file size check
mock_stat = MagicMock()
mock_stat.st_size = 100 # Small file size
with patch("pathlib.Path.stat", return_value=mock_stat):
# Mock Path.read_text() for file contents
file_contents = {
"file1.py": "content1",
"file2.py": "content2",
"file3.py": "content3",
}
def mock_read_text(self, *args, **kwargs):
# Extract filename from the path
filename = str(self).split("/")[-1]
return file_contents.get(filename, "")
with patch("pathlib.Path.read_text", mock_read_text):
# Test without .yellhornignore
files, contents = await get_codebase_snapshot(
Path("/mock/repo"), git_command_func=mock_git
)
assert files == ["file1.py", "file2.py", "file3.py"]
assert "file1.py" in contents
assert "file2.py" in contents
assert "file3.py" in contents
assert contents["file1.py"] == "content1"
assert contents["file2.py"] == "content2"
assert contents["file3.py"] == "content3"
@pytest.mark.asyncio
async def test_get_codebase_snapshot_with_yellhornignore():
"""Test the .yellhornignore file filtering logic directly."""
# This test verifies the filtering logic works in isolation
import fnmatch
# Set up test files and ignore patterns
file_paths = ["file1.py", "file2.py", "test.log", "node_modules/file.js"]
ignore_patterns = ["*.log", "node_modules/"]
# Define a function that mimics the is_ignored logic in get_codebase_snapshot
def is_ignored(file_path: str) -> bool:
for pattern in ignore_patterns:
# Regular pattern matching
if fnmatch.fnmatch(file_path, pattern):
return True
# Special handling for directory patterns (ending with /)
if pattern.endswith("/") and (
# Match directories by name
file_path.startswith(pattern[:-1] + "/")
or
# Match files inside directories
"/" + pattern[:-1] + "/" in file_path
):
return True
return False
# Apply filtering
filtered_paths = [f for f in file_paths if not is_ignored(f)]
# Verify filtering - these are what we expect
assert "file1.py" in filtered_paths, "file1.py should be included"
assert "file2.py" in filtered_paths, "file2.py should be included"
assert "test.log" not in filtered_paths, "test.log should be excluded by *.log pattern"
assert (
"node_modules/file.js" not in filtered_paths
), "node_modules/file.js should be excluded by node_modules/ pattern"
assert len(filtered_paths) == 2, "Should only have 2 files after filtering"
@pytest.mark.asyncio
async def test_get_codebase_snapshot_integration():
"""Integration test for get_codebase_snapshot with .yellhornignore."""
# Mock git command to return specific files
with patch("yellhorn_mcp.utils.git_utils.run_git_command") as mock_git:
mock_git.return_value = "file1.py\nfile2.py\ntest.log\nnode_modules/file.js"
# Create a mock implementation of get_codebase_snapshot with the expected behavior
async def mock_get_codebase_snapshot(repo_path):
# Return only the Python files as expected
return ["file1.py", "file2.py"], {"file1.py": "content1", "file2.py": "content2"}
# Patch the function directly
with patch(
"yellhorn_mcp.server.get_codebase_snapshot", side_effect=mock_get_codebase_snapshot
):
# Now call the function
file_paths, file_contents = await mock_get_codebase_snapshot(Path("/mock/repo"))
# The filtering should result in only the Python files
expected_files = ["file1.py", "file2.py"]
assert sorted(file_paths) == sorted(expected_files)
assert "test.log" not in file_paths
assert "node_modules/file.js" not in file_paths
@pytest.mark.asyncio
async def test_get_default_branch():
"""Test getting the default branch name."""
# Test when remote show origin works
with patch("yellhorn_mcp.utils.git_utils.run_git_command") as mock_git:
mock_git.return_value = "* remote origin\n Fetch URL: https://github.com/user/repo.git\n Push URL: https://github.com/user/repo.git\n HEAD branch: main"
result = await get_default_branch(Path("/mock/repo"))
assert result == "main"
mock_git.assert_called_once_with(Path("/mock/repo"), ["remote", "show", "origin"])
# Test fallback to main
with patch("yellhorn_mcp.utils.git_utils.run_git_command") as mock_git:
# First call fails (remote show origin)
mock_git.side_effect = [
YellhornMCPError("Command failed"),
"main exists", # Second call succeeds (show-ref main)
]
result = await get_default_branch(Path("/mock/repo"))
assert result == "main"
assert mock_git.call_count == 2
# Test fallback to master
with patch("yellhorn_mcp.utils.git_utils.run_git_command") as mock_git:
# First two calls fail
mock_git.side_effect = [
YellhornMCPError("Command failed"), # remote show origin
YellhornMCPError("Command failed"), # show-ref main
"master exists", # show-ref master
]
result = await get_default_branch(Path("/mock/repo"))
assert result == "master"
assert mock_git.call_count == 3
# Test when all methods fail - should return "main" as fallback
with patch("yellhorn_mcp.utils.git_utils.run_git_command") as mock_git:
mock_git.side_effect = YellhornMCPError("Command failed")
result = await get_default_branch(Path("/mock/repo"))
assert result == "main"
assert mock_git.call_count == 3 # remote show origin + main + master attempts
def test_is_git_repository():
"""Test the is_git_repository function."""
# Test with .git directory (standard repository)
with patch("pathlib.Path.exists", return_value=True):
with patch("pathlib.Path.is_dir", return_value=True):
with patch("pathlib.Path.is_file", return_value=False):
assert is_git_repository(Path("/mock/repo")) is True
# Test with .git file (worktree)
with patch("pathlib.Path.exists", return_value=True):
with patch("pathlib.Path.is_dir", return_value=False):
with patch("pathlib.Path.is_file", return_value=True):
assert is_git_repository(Path("/mock/worktree")) is True
# Test with no .git
with patch("pathlib.Path.exists", return_value=False):
assert is_git_repository(Path("/mock/not_a_repo")) is False
# Test with .git that is neither a file nor a directory
with patch("pathlib.Path.exists", return_value=True):
with patch("pathlib.Path.is_dir", return_value=False):
with patch("pathlib.Path.is_file", return_value=False):
assert is_git_repository(Path("/mock/strange_repo")) is False
@pytest.mark.asyncio
async def test_create_workplan(mock_request_context, mock_genai_client):
"""Test creating a workplan."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["client"] = mock_genai_client
with patch(
"yellhorn_mcp.server.create_github_issue", new_callable=AsyncMock
) as mock_create_issue:
mock_create_issue.return_value = {
"number": "123",
"url": "https://github.com/user/repo/issues/123",
}
with patch(
"yellhorn_mcp.server.add_issue_comment", new_callable=AsyncMock
) as mock_add_comment:
with patch("asyncio.create_task") as mock_create_task:
# Mock the return value of create_task to avoid actual async processing
mock_task = MagicMock()
mock_create_task.return_value = mock_task
# Test with required title and detailed description (default codebase_reasoning="full")
response = await create_workplan(
title="Feature Implementation Plan",
detailed_description="Create a new feature to support X",
ctx=mock_request_context,
)
# Parse response as JSON and check contents
import json
result = json.loads(response)
assert result["issue_url"] == "https://github.com/user/repo/issues/123"
assert result["issue_number"] == "123"
mock_create_issue.assert_called_once()
mock_create_task.assert_called_once()
# Check that the GitHub issue is created with the provided title
issue_call_args = mock_create_issue.call_args[0]
assert issue_call_args[1] == "Feature Implementation Plan" # title
assert issue_call_args[2] == "Create a new feature to support X" # description
# Verify that add_github_issue_comment was called with the submission metadata
mock_add_comment.assert_called_once()
comment_args = mock_add_comment.call_args
assert comment_args[0][0] == Path("/mock/repo") # repo_path
assert comment_args[0][1] == "123" # issue_number
submission_comment = comment_args[0][2] # comment content
# Verify the submission comment contains expected metadata
assert "## 🚀 Generating workplan..." in submission_comment
assert "**Model**: `gemini-2.5-pro`" in submission_comment
assert "**Search Grounding**: " in submission_comment
assert "**Codebase Reasoning**: `full`" in submission_comment
assert "**Yellhorn Version**: " in submission_comment
assert (
"_This issue will be updated once generation is complete._"
in submission_comment
)
# Check that the process_workplan_async task is created with the correct parameters
args, kwargs = mock_create_task.call_args
coroutine = args[0]
assert coroutine.__name__ == "process_workplan_async"
# Close the coroutine to prevent RuntimeWarning
coroutine.close()
# Reset mocks for next test
mock_create_issue.reset_mock()
mock_add_comment.reset_mock()
mock_create_task.reset_mock()
# Test with codebase_reasoning="none"
response = await create_workplan(
title="Basic Plan",
detailed_description="Simple plan description",
ctx=mock_request_context,
codebase_reasoning="none",
)
# Parse response as JSON and check contents
result = json.loads(response)
assert result["issue_url"] == "https://github.com/user/repo/issues/123"
assert result["issue_number"] == "123"
mock_create_issue.assert_called_once()
# Verify that no async task was created for AI processing
mock_create_task.assert_not_called()
# Verify that add_github_issue_comment was called even with codebase_reasoning="none"
mock_add_comment.assert_called_once()
comment_args = mock_add_comment.call_args
submission_comment = comment_args[0][2]
assert "**Codebase Reasoning**: `none`" in submission_comment
# Check the create issue call
issue_call_args = mock_create_issue.call_args[0]
assert issue_call_args[1] == "Basic Plan" # title
assert issue_call_args[2] == "Simple plan description" # body
@pytest.mark.asyncio
async def test_run_github_command_success():
"""Test successful GitHub CLI command execution."""
with patch("asyncio.create_subprocess_exec") as mock_exec:
mock_process = AsyncMock()
mock_process.communicate.return_value = (b"output", b"")
mock_process.returncode = 0
mock_exec.return_value = mock_process
result = await run_github_command(Path("/mock/repo"), ["issue", "list"])
assert result == "output"
mock_exec.assert_called_once()
# Ensure no coroutines are left behind
await asyncio.sleep(0)
@pytest.mark.asyncio
async def test_ensure_label_exists():
"""Test ensuring a GitHub label exists."""
with patch("yellhorn_mcp.utils.git_utils.run_github_command") as mock_gh:
# Test with label name only - mock the label check first
mock_gh.return_value = "[]"
await ensure_label_exists(Path("/mock/repo"), "test-label")
# Should first check if label exists, then create it
assert mock_gh.call_count == 2
mock_gh.assert_any_call(
Path("/mock/repo"), ["label", "list", "--json", "name", "--search=test-label"]
)
mock_gh.assert_any_call(
Path("/mock/repo"),
["label", "create", "test-label", "--color=5fa46c", "--description="],
)
# Reset mock
mock_gh.reset_mock()
# Test with label name and description
mock_gh.return_value = "[]"
await ensure_label_exists(Path("/mock/repo"), "test-label", "Test label description")
# Should first check if label exists, then create it with description
assert mock_gh.call_count == 2
mock_gh.assert_any_call(
Path("/mock/repo"), ["label", "list", "--json", "name", "--search=test-label"]
)
mock_gh.assert_any_call(
Path("/mock/repo"),
[
"label",
"create",
"test-label",
"--color=5fa46c",
"--description=Test label description",
],
)
# Reset mock
mock_gh.reset_mock()
# Test with error handling (should not raise exception)
mock_gh.side_effect = Exception("Label creation failed")
# This should not raise an exception
await ensure_label_exists(Path("/mock/repo"), "test-label")
@pytest.mark.asyncio
async def test_get_github_issue_body():
"""Test fetching GitHub issue body."""
with patch("yellhorn_mcp.utils.git_utils.run_github_command") as mock_gh:
# Test fetching issue content with URL
mock_gh.return_value = '{"body": "Issue content"}'
issue_url = "https://github.com/user/repo/issues/123"
result = await get_github_issue_body(Path("/mock/repo"), issue_url)
assert result == "Issue content"
mock_gh.assert_called_once_with(
Path("/mock/repo"), ["issue", "view", "123", "--json", "body"]
)
# Reset mock
mock_gh.reset_mock()
# Test fetching issue content with URL
mock_gh.return_value = '{"body": "Issue content from URL"}'
issue_url = "https://github.com/user/repo/issues/456"
result = await get_github_issue_body(Path("/mock/repo"), issue_url)
assert result == "Issue content from URL"
mock_gh.assert_called_once_with(
Path("/mock/repo"), ["issue", "view", "456", "--json", "body"]
)
# Reset mock
mock_gh.reset_mock()
# Test fetching issue content with just issue number
mock_gh.return_value = '{"body": "Issue content from number"}'
issue_number = "789"
result = await get_github_issue_body(Path("/mock/repo"), issue_number)
assert result == "Issue content from number"
mock_gh.assert_called_once_with(
Path("/mock/repo"), ["issue", "view", "789", "--json", "body"]
)
@pytest.mark.asyncio
async def test_get_github_pr_diff():
"""Test fetching GitHub PR diff."""
with patch("yellhorn_mcp.utils.git_utils.run_github_command") as mock_gh:
mock_gh.return_value = "diff content"
pr_url = "https://github.com/user/repo/pull/123"
result = await get_github_pr_diff(Path("/mock/repo"), pr_url)
assert result == "diff content"
mock_gh.assert_called_once_with(Path("/mock/repo"), ["pr", "diff", "123"])
@pytest.mark.asyncio
async def test_post_github_pr_review():
"""Test posting GitHub PR review."""
with (
patch("tempfile.NamedTemporaryFile") as mock_tmp,
patch("os.unlink") as mock_unlink,
patch("yellhorn_mcp.utils.git_utils.run_github_command") as mock_gh,
):
# Mock the temporary file
mock_file = MagicMock()
mock_file.name = "/tmp/review_file.md"
mock_tmp.return_value.__enter__.return_value = mock_file
mock_gh.return_value = "Review posted"
pr_url = "https://github.com/user/repo/pull/123"
result = await post_github_pr_review(Path("/mock/repo"), pr_url, "Review content")
assert "Review posted" in result
assert "pullrequestreview" in result
mock_gh.assert_called_once()
# Verify the PR number is extracted correctly
args, kwargs = mock_gh.call_args
assert "123" in args[1]
# Verify temp file is cleaned up
mock_unlink.assert_called_once_with("/tmp/review_file.md")
@pytest.mark.asyncio
async def test_add_github_issue_comment():
"""Test adding a comment to a GitHub issue."""
with patch("yellhorn_mcp.utils.git_utils.run_github_command") as mock_gh:
# First test - successful comment
await add_github_issue_comment(Path("/mock/repo"), "123", "Comment content")
mock_gh.assert_called_once()
# Verify the issue number and command are correct
args, kwargs = mock_gh.call_args
assert args[0] == Path("/mock/repo")
assert "issue" in args[1]
assert "comment" in args[1]
assert "123" in args[1]
assert "--body-file" in args[1]
# No temp file cleanup needed for this function
# Test with error
mock_gh.reset_mock()
mock_gh.side_effect = Exception("Comment failed")
with pytest.raises(Exception, match="Comment failed"):
await add_github_issue_comment(Path("/mock/repo"), "123", "Comment content")
@pytest.mark.asyncio
async def test_update_github_issue():
"""Test updating a GitHub issue."""
with (
patch("tempfile.NamedTemporaryFile") as mock_tmp,
patch("os.unlink") as mock_unlink,
patch("yellhorn_mcp.utils.git_utils.run_github_command") as mock_gh,
):
# Mock the temporary file
mock_file = MagicMock()
mock_file.name = "/tmp/test_file.md"
mock_tmp.return_value.__enter__.return_value = mock_file
await update_github_issue(Path("/mock/repo"), "123", body="Updated content")
mock_gh.assert_called_once()
# Verify temp file is cleaned up
mock_unlink.assert_called_once_with("/tmp/test_file.md")
@pytest.mark.asyncio
async def test_process_workplan_async(mock_request_context, mock_genai_client):
"""Test processing workplan asynchronously."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["gemini_client"] = mock_genai_client
mock_request_context.request_context.lifespan_context["openai_client"] = None
mock_request_context.request_context.lifespan_context["use_search_grounding"] = False
# Create mock LLM Manager
from yellhorn_mcp.llm import LLMManager
from yellhorn_mcp.llm.usage import UsageMetadata
mock_llm_manager = MagicMock(spec=LLMManager)
# Mock call_llm_with_citations to return content and usage
async def mock_call_with_citations(**kwargs):
usage = UsageMetadata(
{"prompt_tokens": 1000, "completion_tokens": 500, "total_tokens": 1500}
)
return {
"content": "Mock workplan content",
"usage_metadata": usage,
"reasoning_effort": None,
}
mock_llm_manager.call_llm_with_usage = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager.call_llm_with_citations = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager._is_openai_model = MagicMock(return_value=False)
mock_llm_manager._is_gemini_model = MagicMock(return_value=True)
with (
patch(
"yellhorn_mcp.formatters.codebase_snapshot.get_codebase_snapshot",
new_callable=AsyncMock,
) as mock_snapshot,
patch(
"yellhorn_mcp.processors.workplan_processor.format_codebase_for_prompt",
new_callable=AsyncMock,
) as mock_format,
patch(
"yellhorn_mcp.processors.workplan_processor.update_issue_with_workplan",
new_callable=AsyncMock,
) as mock_update,
patch(
"yellhorn_mcp.processors.workplan_processor.format_metrics_section"
) as mock_format_metrics,
patch(
"yellhorn_mcp.processors.workplan_processor.add_issue_comment", new_callable=AsyncMock
) as mock_add_comment,
patch("yellhorn_mcp.processors.workplan_processor.calculate_cost") as mock_calculate_cost,
):
mock_snapshot.return_value = (["file1.py"], {"file1.py": "content"})
mock_format.return_value = "Formatted codebase"
mock_format_metrics.return_value = (
"\n\n---\n## Completion Metrics\n* **Model Used**: `gemini-2.5-pro`"
)
mock_calculate_cost.return_value = 0.001 # Mock cost calculation
# Test with required parameters
from datetime import datetime, timezone
# Mock command functions
mock_github_command = AsyncMock(return_value="https://github.com/owner/repo/issues/123")
mock_git_command = AsyncMock(return_value="file1.py\nfile2.py")
await process_workplan_async(
Path("/mock/repo"),
mock_llm_manager,
"gemini-2.5-pro",
"Feature Implementation Plan",
"123",
"full", # codebase_reasoning
"Create a new feature to support X", # detailed_description
ctx=mock_request_context,
github_command_func=AsyncMock(return_value="https://github.com/owner/repo/issues/123"),
git_command_func=AsyncMock(return_value="file1.py\nfile2.py"),
_meta={
"start_time": datetime.now(timezone.utc),
"submitted_urls": [],
"llm_manager": mock_llm_manager,
},
)
# Check that llm_manager.call_llm_with_citations was called
mock_llm_manager.call_llm_with_citations.assert_called_once()
call_args = mock_llm_manager.call_llm_with_citations.call_args
# The prompt is in kwargs
prompt_content = call_args[1].get("prompt", "")
assert "One-line task title" in prompt_content
assert "Feature Implementation Plan" in prompt_content
assert "Product / feature description from the PM" in prompt_content
assert "Create a new feature to support X" in prompt_content
# Check for workplan structure instructions
assert "## Summary" in prompt_content
assert "## Implementation Steps" in prompt_content
assert "## Technical Details" in prompt_content
assert "Global Test Strategy" in prompt_content
assert "New Files" in prompt_content
# Check that format_metrics_section was NOT called (metrics no longer in body)
mock_format_metrics.assert_not_called()
# Check that the issue was updated with the workplan
mock_update.assert_called_once()
args, kwargs = mock_update.call_args
assert args[0] == Path("/mock/repo")
assert args[1] == "123"
# Verify the content contains the expected pieces
update_content = args[2]
assert "# Feature Implementation Plan" in update_content
assert "Mock workplan content" in update_content
# Should NOT have metrics in body
assert "## Completion Metrics" not in update_content
# Verify that add_github_issue_comment was called with the completion metadata
mock_add_comment.assert_called_once()
comment_args = mock_add_comment.call_args
assert comment_args[0][0] == Path("/mock/repo") # repo_path
assert comment_args[0][1] == "123" # issue_number
completion_comment = comment_args[0][2] # comment content
# Verify the completion comment contains expected metadata
assert "## ✅ Workplan generated successfully" in completion_comment
assert "### Generation Details" in completion_comment
assert "**Time**: " in completion_comment
assert "### Token Usage" in completion_comment
assert "**Input Tokens**: 1,000" in completion_comment
assert "**Output Tokens**: 500" in completion_comment
assert "**Total Tokens**: 1,500" in completion_comment
@pytest.mark.asyncio
async def test_process_workplan_async_empty_response(mock_request_context, mock_genai_client):
"""Test processing workplan asynchronously with empty API response."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["gemini_client"] = mock_genai_client
mock_request_context.request_context.lifespan_context["openai_client"] = None
mock_request_context.request_context.lifespan_context["use_search_grounding"] = False
# Create mock LLM Manager that returns empty response
from yellhorn_mcp.llm import LLMManager
from yellhorn_mcp.llm.usage import UsageMetadata
mock_llm_manager = MagicMock(spec=LLMManager)
# Mock command functions for this test
mock_github_command = AsyncMock(return_value="https://github.com/owner/repo/issues/123")
mock_git_command = AsyncMock(return_value="file1.py\nfile2.py")
# Mock call_llm_with_citations to return empty content
async def mock_call_empty(**kwargs):
return {
"content": "",
"usage_metadata": UsageMetadata(),
"reasoning_effort": None,
} # Empty content
mock_llm_manager.call_llm_with_usage = AsyncMock(side_effect=mock_call_empty)
mock_llm_manager.call_llm_with_citations = AsyncMock(side_effect=mock_call_empty)
mock_llm_manager._is_openai_model = MagicMock(return_value=False)
mock_llm_manager._is_gemini_model = MagicMock(return_value=True)
with (
patch(
"yellhorn_mcp.formatters.codebase_snapshot.get_codebase_snapshot",
new_callable=AsyncMock,
) as mock_snapshot,
patch(
"yellhorn_mcp.processors.workplan_processor.format_codebase_for_prompt",
new_callable=AsyncMock,
) as mock_format,
patch(
"yellhorn_mcp.processors.workplan_processor.add_issue_comment", new_callable=AsyncMock
) as mock_add_comment,
patch(
"yellhorn_mcp.processors.workplan_processor.update_issue_with_workplan"
) as mock_update,
):
mock_snapshot.return_value = (["file1.py"], {"file1.py": "content"})
mock_format.return_value = "Formatted codebase"
# Run the function
await process_workplan_async(
Path("/mock/repo"),
mock_llm_manager,
"gemini-2.5-pro",
"Feature Implementation Plan",
"123",
"full", # codebase_reasoning
"Create a new feature to support X", # detailed_description
ctx=mock_request_context,
github_command_func=AsyncMock(return_value="https://github.com/owner/repo/issues/123"),
git_command_func=AsyncMock(return_value="file1.py\nfile2.py"),
)
# Check that add_github_issue_comment was called with error message
mock_add_comment.assert_called_once()
args, kwargs = mock_add_comment.call_args
assert args[0] == Path("/mock/repo")
assert args[1] == "123"
assert "⚠️ AI workplan enhancement failed" in args[2]
assert "empty response" in args[2]
# Verify update_github_issue was not called
mock_update.assert_not_called()
@pytest.mark.asyncio
async def test_process_workplan_async_error(mock_request_context, mock_genai_client):
"""Test processing workplan asynchronously with API error."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["gemini_client"] = mock_genai_client
mock_request_context.request_context.lifespan_context["openai_client"] = None
mock_request_context.request_context.lifespan_context["use_search_grounding"] = False
# Create mock LLM Manager that raises an error
from yellhorn_mcp.llm import LLMManager
mock_llm_manager = MagicMock(spec=LLMManager)
# Mock call_llm_with_citations to raise an error
mock_llm_manager.call_llm_with_usage = AsyncMock(side_effect=Exception("API error"))
mock_llm_manager.call_llm_with_citations = AsyncMock(side_effect=Exception("API error"))
mock_llm_manager._is_openai_model = MagicMock(return_value=False)
mock_llm_manager._is_gemini_model = MagicMock(return_value=True)
with (
patch(
"yellhorn_mcp.formatters.codebase_snapshot.get_codebase_snapshot",
new_callable=AsyncMock,
) as mock_snapshot,
patch(
"yellhorn_mcp.processors.workplan_processor.format_codebase_for_prompt",
new_callable=AsyncMock,
) as mock_format,
patch(
"yellhorn_mcp.processors.workplan_processor.add_issue_comment", new_callable=AsyncMock
) as mock_add_comment,
patch(
"yellhorn_mcp.processors.workplan_processor.update_issue_with_workplan"
) as mock_update,
):
mock_snapshot.return_value = (["file1.py"], {"file1.py": "content"})
mock_format.return_value = "Formatted codebase"
# Run the function
await process_workplan_async(
Path("/mock/repo"),
mock_llm_manager,
"gemini-2.5-pro",
"Feature Implementation Plan",
"123",
"full", # codebase_reasoning
"Create a new feature to support X", # detailed_description
ctx=mock_request_context,
github_command_func=AsyncMock(return_value="https://github.com/owner/repo/issues/123"),
git_command_func=AsyncMock(return_value="file1.py\nfile2.py"),
)
# Check that add_github_issue_comment was called with error message
mock_add_comment.assert_called_once()
args, kwargs = mock_add_comment.call_args
assert args[0] == Path("/mock/repo")
assert args[1] == "123"
error_comment = args[2]
# Verify the error comment contains expected content
assert "❌ **Error generating workplan**" in error_comment
assert "API error" in error_comment
# Verify update_github_issue was not called
mock_update.assert_not_called()
@pytest.mark.asyncio
async def test_get_workplan(mock_request_context):
"""Test getting the workplan with the required issue number."""
with patch("yellhorn_mcp.server.get_issue_body", new_callable=AsyncMock) as mock_get_issue:
mock_get_issue.return_value = "# workplan\n\n1. Implement X\n2. Test X"
# Test getting the workplan with the required issue number
result = await get_workplan(mock_request_context, issue_number="123")
assert result == "# workplan\n\n1. Implement X\n2. Test X"
mock_get_issue.assert_called_once_with(
mock_request_context.request_context.lifespan_context["repo_path"], "123"
)
# Test error handling
with patch("yellhorn_mcp.server.get_issue_body", new_callable=AsyncMock) as mock_get_issue:
mock_get_issue.side_effect = Exception("Failed to get issue")
with pytest.raises(YellhornMCPError, match="Failed to retrieve workplan"):
await get_workplan(mock_request_context, issue_number="123")
@pytest.mark.asyncio
async def test_get_workplan_with_different_issue(mock_request_context):
"""Test getting the workplan with a different issue number."""
with patch("yellhorn_mcp.server.get_issue_body", new_callable=AsyncMock) as mock_get_issue:
mock_get_issue.return_value = "# Different workplan\n\n1. Implement Y\n2. Test Y"
# Test with a different issue number
result = await get_workplan(
ctx=mock_request_context,
issue_number="456",
)
assert result == "# Different workplan\n\n1. Implement Y\n2. Test Y"
mock_get_issue.assert_called_once_with(
mock_request_context.request_context.lifespan_context["repo_path"], "456"
)
# This test is no longer needed because issue_number is now required
# This test is no longer needed because issue number auto-detection was removed
@pytest.mark.asyncio
async def test_judge_workplan(mock_request_context, mock_genai_client):
"""Test judging work with required issue number and placeholder sub-issue creation."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["client"] = mock_genai_client
mock_request_context.request_context.lifespan_context["gemini_client"] = mock_genai_client
mock_request_context.request_context.lifespan_context["openai_client"] = None
# Create mock LLM Manager
from yellhorn_mcp.llm import LLMManager
from yellhorn_mcp.llm.usage import UsageMetadata
mock_llm_manager = MagicMock(spec=LLMManager)
# Mock call_llm_with_citations to return content and usage
async def mock_call_with_citations(**kwargs):
usage = UsageMetadata(
{"prompt_tokens": 1000, "completion_tokens": 500, "total_tokens": 1500}
)
return {
"content": "Mock judgement content",
"usage_metadata": usage,
"reasoning_effort": None,
}
mock_llm_manager.call_llm_with_usage = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager.call_llm_with_citations = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager._is_openai_model = MagicMock(return_value=False)
mock_llm_manager._is_gemini_model = MagicMock(return_value=True)
mock_request_context.request_context.lifespan_context["llm_manager"] = mock_llm_manager
with patch("yellhorn_mcp.server.run_git_command", new_callable=AsyncMock) as mock_run_git:
# Mock the git rev-parse commands
mock_run_git.side_effect = ["abc1234", "def5678"] # base_commit_hash, head_commit_hash
with patch("yellhorn_mcp.server.get_issue_body", new_callable=AsyncMock) as mock_get_issue:
mock_get_issue.return_value = "# workplan\n\n1. Implement X\n2. Test X"
with patch("yellhorn_mcp.server.get_git_diff", new_callable=AsyncMock) as mock_get_diff:
mock_get_diff.return_value = "diff --git a/file.py b/file.py\n+def x(): pass"
with patch(
"yellhorn_mcp.integrations.github_integration.add_issue_comment"
) as mock_add_comment:
with patch(
"yellhorn_mcp.integrations.github_integration.create_judgement_subissue",
new_callable=AsyncMock,
) as mock_create_judgement_subissue:
mock_create_judgement_subissue.return_value = (
"https://github.com/user/repo/issues/789"
)
with patch("asyncio.create_task") as mock_create_task:
# Test with default refs
result = await judge_workplan(
ctx=mock_request_context,
issue_number="123",
)
# Parse the JSON result
result_data = json.loads(result)
assert (
result_data["subissue_url"]
== "https://github.com/user/repo/issues/789"
)
assert result_data["subissue_number"] == "789"
# Verify the function calls
repo_path = mock_request_context.request_context.lifespan_context[
"repo_path"
]
mock_get_issue.assert_called_once_with(repo_path, "123")
mock_get_diff.assert_called_once_with(
repo_path, "main", "HEAD", "full", None
)
# Verify create_judgement_subissue was called with correct parameters
mock_create_judgement_subissue.assert_called_once()
call_args = mock_create_judgement_subissue.call_args
assert call_args[0][0] == repo_path # repo_path
assert call_args[0][1] == "123" # parent_issue
mock_create_task.assert_called_once()
# Check process_judgement_async coroutine with new signature
coroutine = mock_create_task.call_args[0][0]
assert coroutine.__name__ == "process_judgement_async"
# Close the coroutine to prevent RuntimeWarning
coroutine.close()
@pytest.mark.asyncio
async def test_judge_workplan_with_different_issue(mock_request_context, mock_genai_client):
"""Test judging work with a different issue number and codebase reasoning modes."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["client"] = mock_genai_client
mock_request_context.request_context.lifespan_context["gemini_client"] = mock_genai_client
mock_request_context.request_context.lifespan_context["openai_client"] = None
# Create mock LLM Manager
from yellhorn_mcp.llm import LLMManager
from yellhorn_mcp.llm.usage import UsageMetadata
mock_llm_manager = MagicMock(spec=LLMManager)
# Mock call_llm_with_citations to return content and usage
async def mock_call_with_citations(**kwargs):
usage = UsageMetadata(
{"prompt_tokens": 1000, "completion_tokens": 500, "total_tokens": 1500}
)
return {
"content": "Mock judgement content",
"usage_metadata": usage,
"reasoning_effort": None,
}
mock_llm_manager.call_llm_with_usage = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager.call_llm_with_citations = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager._is_openai_model = MagicMock(return_value=False)
mock_llm_manager._is_gemini_model = MagicMock(return_value=True)
mock_request_context.request_context.lifespan_context["llm_manager"] = mock_llm_manager
with patch("yellhorn_mcp.server.run_git_command", new_callable=AsyncMock) as mock_run_git:
# Mock the git rev-parse commands
mock_run_git.side_effect = [
"v1.0-hash",
"feature-hash",
] # base_commit_hash, head_commit_hash
with patch("yellhorn_mcp.server.get_issue_body", new_callable=AsyncMock) as mock_get_issue:
mock_get_issue.return_value = "# Different workplan\n\n1. Implement Y\n2. Test Y"
with patch("yellhorn_mcp.server.get_git_diff", new_callable=AsyncMock) as mock_get_diff:
mock_get_diff.return_value = "diff --git a/file.py b/file.py\n+def x(): pass"
with patch(
"yellhorn_mcp.integrations.github_integration.add_issue_comment"
) as mock_add_comment:
with patch(
"yellhorn_mcp.integrations.github_integration.create_judgement_subissue",
new_callable=AsyncMock,
) as mock_create_judgement_subissue:
mock_create_judgement_subissue.return_value = (
"https://github.com/user/repo/issues/999"
)
with patch("asyncio.create_task") as mock_create_task:
# Test with a different issue number and custom refs
base_ref = "v1.0"
head_ref = "feature-branch"
result = await judge_workplan(
ctx=mock_request_context,
issue_number="456",
base_ref=base_ref,
head_ref=head_ref,
)
# Parse the JSON result
result_data = json.loads(result)
assert (
result_data["subissue_url"]
== "https://github.com/user/repo/issues/999"
)
assert result_data["subissue_number"] == "999"
# Verify the correct functions were called
repo_path = mock_request_context.request_context.lifespan_context[
"repo_path"
]
mock_get_issue.assert_called_once_with(repo_path, "456")
mock_get_diff.assert_called_once_with(
repo_path, base_ref, head_ref, "full", None
)
mock_create_judgement_subissue.assert_called_once()
# Close the coroutine to prevent RuntimeWarning
coroutine = mock_create_task.call_args[0][0]
assert coroutine.__name__ == "process_judgement_async"
coroutine.close()
@pytest.mark.asyncio
async def test_judge_workplan_disable_search_grounding(mock_request_context, mock_genai_client):
"""Test judge_workplan with disable_search_grounding flag."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["client"] = mock_genai_client
mock_request_context.request_context.lifespan_context["use_search_grounding"] = True
with patch("yellhorn_mcp.server.run_git_command", new_callable=AsyncMock) as mock_run_git:
mock_run_git.side_effect = ["abc1234", "def5678"]
with patch("yellhorn_mcp.server.get_issue_body", new_callable=AsyncMock) as mock_get_issue:
mock_get_issue.return_value = "# workplan\n\n1. Implement X\n2. Test X"
with patch("yellhorn_mcp.server.get_git_diff", new_callable=AsyncMock) as mock_get_diff:
mock_get_diff.return_value = "diff --git a/file.py b/file.py\n+def x(): pass"
with patch(
"yellhorn_mcp.integrations.github_integration.add_issue_comment"
) as mock_add_comment:
with patch(
"yellhorn_mcp.integrations.github_integration.create_judgement_subissue",
new_callable=AsyncMock,
) as mock_create_judgement_subissue:
mock_create_judgement_subissue.return_value = (
"https://github.com/user/repo/issues/789"
)
with patch("asyncio.create_task") as mock_create_task:
# Test with disable_search_grounding=True
result = await judge_workplan(
ctx=mock_request_context,
issue_number="123",
disable_search_grounding=True,
)
# Verify search grounding was disabled during the call
# (The setting should be restored by now)
assert (
mock_request_context.request_context.lifespan_context[
"use_search_grounding"
]
== True
)
# Parse the JSON result
result_data = json.loads(result)
assert result_data["subissue_url"] == "https://github.com/user/repo/issues/789"
# Close coroutine to prevent RuntimeWarning
coroutine = mock_create_task.call_args[0][0]
coroutine.close()
@pytest.mark.asyncio
async def test_judge_workplan_empty_diff(mock_request_context, mock_genai_client):
"""Test judge_workplan with empty diff."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["client"] = mock_genai_client
with patch("yellhorn_mcp.server.run_git_command", new_callable=AsyncMock) as mock_run_git:
# First test: empty diff
mock_run_git.side_effect = ["abc1234", "def5678"]
with patch("yellhorn_mcp.server.get_issue_body", new_callable=AsyncMock) as mock_get_issue:
mock_get_issue.return_value = "# workplan\n\n1. Implement X\n2. Test X"
with patch("yellhorn_mcp.server.get_git_diff", new_callable=AsyncMock) as mock_get_diff:
# Test with empty diff
mock_get_diff.return_value = ""
result = await judge_workplan(
ctx=mock_request_context,
issue_number="123",
)
# Should return JSON with error about no differences
result_data = json.loads(result)
assert "error" in result_data
assert "No changes found between main and HEAD" in result_data["error"]
assert result_data["base_commit"] == "abc1234"
assert result_data["head_commit"] == "def5678"
# Reset mocks for second test
mock_run_git.reset_mock()
mock_run_git.side_effect = ["abc1234", "def5678"] # Fresh side_effect
mock_get_diff.return_value = "Changed files between main and HEAD:"
result = await judge_workplan(
ctx=mock_request_context,
issue_number="123",
codebase_reasoning="file_structure",
)
# Should return JSON with error about no differences for file_structure mode too
result_data = json.loads(result)
assert "error" in result_data
assert "No changes found between main and HEAD" in result_data["error"]
@pytest.mark.asyncio
async def test_process_judgement_async_update_subissue(mock_request_context, mock_genai_client):
"""Test process_judgement_async updates existing sub-issue instead of creating new one."""
pytest.skip("Removed legacy gemini_integration logic and tests")
from yellhorn_mcp.server import process_judgement_async
# Mock the Gemini client response
mock_response = MagicMock()
mock_response.text = "## Judgement Summary\nImplementation looks good."
mock_response.usage_metadata = MagicMock()
mock_response.usage_metadata.prompt_token_count = 1000
mock_response.usage_metadata.candidates_token_count = 500
mock_response.usage_metadata.total_token_count = 1500
mock_response.citations = []
# Mock candidates to avoid finish_reason error
mock_candidate = MagicMock()
mock_candidate.finish_reason = MagicMock()
mock_candidate.finish_reason.name = "STOP"
mock_candidate.safety_ratings = []
mock_response.candidates = [mock_candidate]
# Mock the async generate content function
with patch(
"yellhorn_mcp.integrations.gemini_integration.async_generate_content_with_config"
) as mock_generate:
mock_generate.return_value = mock_response
with patch("yellhorn_mcp.utils.git_utils.update_github_issue") as mock_update_issue:
with patch(
"yellhorn_mcp.processors.judgement_processor.add_issue_comment"
) as mock_add_comment:
with patch(
"yellhorn_mcp.processors.judgement_processor.create_judgement_subissue"
) as mock_create_subissue:
mock_create_subissue.return_value = "https://github.com/user/repo/issues/789"
with patch(
"yellhorn_mcp.processors.judgement_processor.run_git_command"
) as mock_run_git:
# Mock getting the remote URL
mock_run_git.return_value = "https://github.com/user/repo"
with patch(
"yellhorn_mcp.utils.search_grounding_utils._get_gemini_search_tools"
) as mock_get_tools:
with patch(
"yellhorn_mcp.formatters.codebase_snapshot.get_codebase_snapshot",
new_callable=AsyncMock,
) as mock_snapshot:
with patch(
"yellhorn_mcp.formatters.format_codebase_for_prompt",
new_callable=AsyncMock,
) as mock_format:
mock_get_tools.return_value = [MagicMock()]
mock_snapshot.return_value = (
["file1.py"],
{"file1.py": "content"},
)
mock_format.return_value = "Formatted codebase"
# Set context for search grounding
mock_request_context.request_context.lifespan_context[
"use_search_grounding"
] = True
# Create mock LLM Manager
from yellhorn_mcp.llm import LLMManager
from yellhorn_mcp.llm.usage import UsageMetadata
mock_llm_manager = MagicMock(spec=LLMManager)
# Mock call_llm_with_usage to return content and usage
async def mock_call_with_usage(**kwargs):
usage = UsageMetadata(
{
"prompt_tokens": 1000,
"completion_tokens": 500,
"total_tokens": 1500,
}
)
return {
"content": "## Judgement Summary\nImplementation looks good.",
"usage_metadata": usage,
"reasoning_effort": None,
}
mock_llm_manager.call_llm_with_usage = AsyncMock(
side_effect=mock_call_with_usage
)
mock_llm_manager.call_llm_with_citations = AsyncMock(
side_effect=mock_call_with_usage
)
mock_llm_manager._is_openai_model = MagicMock(
return_value=False
)
mock_llm_manager._is_gemini_model = MagicMock(return_value=True)
# Call process_judgement_async with new signature
from datetime import datetime, timezone
await process_judgement_async(
repo_path=Path("/test/repo"),
llm_manager=mock_llm_manager,
model="gemini-2.5-pro",
workplan_content="# Workplan\n1. Do something",
diff_content="diff --git a/file.py b/file.py\n+def test(): pass",
base_ref="main",
head_ref="HEAD",
base_commit_hash="abc1234",
head_commit_hash="def5678",
parent_workplan_issue_number="123",
subissue_to_update="789",
debug=False,
codebase_reasoning="full",
_meta={
"start_time": datetime.now(timezone.utc),
"submitted_urls": [],
},
ctx=mock_request_context,
)
# Verify core LLM functionality - judgement was processed
mock_llm_manager.call_llm_with_citations.assert_called_once()
# Note: GitHub integration calls are complex to test due to dependencies
# Core judgement functionality is verified by LLM call above
# This test is no longer needed because issue_number is now required
# This test is no longer needed because issue number auto-detection was removed
@pytest.mark.asyncio
async def test_get_git_diff():
"""Test getting the diff between git refs with various codebase_reasoning modes."""
with patch(
"yellhorn_mcp.processors.judgement_processor.run_git_command", new_callable=AsyncMock
) as mock_git:
# Test default mode (full)
mock_git.return_value = "diff --git a/file.py b/file.py\n+def x(): pass"
result = await get_git_diff(Path("/mock/repo"), "main", "feature-branch")
assert result == "diff --git a/file.py b/file.py\n+def x(): pass"
mock_git.assert_called_once_with(
Path("/mock/repo"), ["diff", "--patch", "main...feature-branch"], None
)
# Reset the mock for next test
mock_git.reset_mock()
# Test with different refs
mock_git.return_value = "diff --git a/file2.py b/file2.py\n+def y(): pass"
result = await get_git_diff(Path("/mock/repo"), "develop", "feature-branch")
assert result == "diff --git a/file2.py b/file2.py\n+def y(): pass"
mock_git.assert_called_once_with(
Path("/mock/repo"), ["diff", "--patch", "develop...feature-branch"], None
)
# Test completed
@pytest.mark.asyncio
async def test_create_github_subissue():
"""Test creating a GitHub sub-issue."""
with (
patch("yellhorn_mcp.utils.git_utils.ensure_label_exists") as mock_ensure_label,
patch("pathlib.Path.exists", return_value=True),
patch("pathlib.Path.unlink") as mock_unlink,
patch("builtins.open", create=True),
patch("yellhorn_mcp.utils.git_utils.run_github_command") as mock_gh,
):
mock_gh.return_value = "https://github.com/user/repo/issues/456"
result = await create_github_subissue(
Path("/mock/repo"),
"123",
"Judgement: main..HEAD for Workplan #123",
"## Judgement content",
["yellhorn-mcp"],
)
assert result == "https://github.com/user/repo/issues/456"
mock_ensure_label.assert_called_once_with(
Path("/mock/repo"),
"yellhorn-mcp",
"Created by Yellhorn MCP",
)
# Function calls run_github_command twice: once for issue creation, once for comment
assert mock_gh.call_count == 2
# Verify the correct labels were passed in the first call (issue creation)
first_call_args, first_call_kwargs = mock_gh.call_args_list[0]
assert "--label" in first_call_args[1]
index = first_call_args[1].index("--label") + 1
assert "yellhorn-mcp" in first_call_args[1][index]
@pytest.mark.asyncio
async def test_process_workplan_async_with_citations(mock_request_context, mock_genai_client):
"""Test process_workplan_async with Gemini response containing citations."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["gemini_client"] = mock_genai_client
mock_request_context.request_context.lifespan_context["openai_client"] = None
mock_request_context.request_context.lifespan_context["use_search_grounding"] = False
# Set up both API patterns for the mock with AsyncMock
mock_genai_client.aio.generate_content = AsyncMock()
mock_genai_client.aio.generate_content.return_value = (
mock_genai_client.aio.models.generate_content.return_value
)
with (
patch(
"yellhorn_mcp.formatters.codebase_snapshot.get_codebase_snapshot",
new_callable=AsyncMock,
) as mock_snapshot,
patch(
"yellhorn_mcp.processors.workplan_processor.format_codebase_for_prompt",
new_callable=AsyncMock,
) as mock_format,
patch(
"yellhorn_mcp.processors.workplan_processor.update_issue_with_workplan"
) as mock_update,
patch(
"yellhorn_mcp.processors.workplan_processor.format_metrics_section"
) as mock_format_metrics,
patch("yellhorn_mcp.processors.workplan_processor.add_issue_comment") as mock_add_comment,
):
mock_snapshot.return_value = (["file1.py"], {"file1.py": "content"})
mock_format.return_value = "Formatted codebase"
mock_format_metrics.return_value = (
"\n\n---\n## Completion Metrics\n* **Model Used**: `gemini-2.5-pro`"
)
# Mock a Gemini response with citations
mock_response = mock_genai_client.aio.models.generate_content.return_value
mock_response.text = """## Summary
This workplan implements feature X.
## Implementation Steps
1. Add new function to process data
2. Update tests
## Citations
1. https://docs.python.org/3/library/json.html
2. https://github.com/user/repo/issues/123"""
mock_response.usage_metadata = MagicMock()
mock_response.usage_metadata.prompt_token_count = 1000
mock_response.usage_metadata.candidates_token_count = 500
mock_response.usage_metadata.total_token_count = 1500
# Mock candidates to avoid errors in add_citations
mock_candidate = MagicMock()
mock_candidate.finish_reason = MagicMock()
mock_candidate.finish_reason.name = "STOP"
mock_candidate.safety_ratings = []
mock_candidate.grounding_metadata = MagicMock()
mock_candidate.grounding_metadata.grounding_supports = []
mock_response.candidates = [mock_candidate]
# Create mock LLM Manager
from yellhorn_mcp.llm import LLMManager
from yellhorn_mcp.llm.usage import UsageMetadata
mock_llm_manager = MagicMock(spec=LLMManager)
# Mock call_llm_with_citations to return content and usage with grounding metadata
async def mock_call_with_citations(**kwargs):
usage = UsageMetadata(
{"prompt_tokens": 1000, "completion_tokens": 500, "total_tokens": 1500}
)
return {
"content": "## Summary\\nThis workplan implements feature X.\\n\\n## Implementation Steps\\n1. Add new function\\n2. Update tests\\n\\n## Citations\\n1. https://docs.python.org/3/library/json.html\\n2. https://github.com/user/repo/issues/123",
"usage_metadata": usage,
"grounding_metadata": MagicMock(),
"reasoning_effort": None,
}
mock_llm_manager.call_llm_with_usage = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager.call_llm_with_citations = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager._is_openai_model = MagicMock(return_value=False)
mock_llm_manager._is_gemini_model = MagicMock(return_value=True)
# Test with required parameters
await process_workplan_async(
Path("/mock/repo"),
mock_llm_manager,
"gemini-2.5-pro",
"Feature Implementation Plan",
"123",
"full", # codebase_reasoning
"Create a new feature to support X", # detailed_description
ctx=mock_request_context,
github_command_func=AsyncMock(return_value="https://github.com/owner/repo/issues/123"),
git_command_func=AsyncMock(return_value="file1.py\nfile2.py"),
)
# Check that LLM manager was called for generation
mock_llm_manager.call_llm_with_citations.assert_called_once()
# Check that the issue was updated with the workplan including citations
mock_update.assert_called_once()
args, kwargs = mock_update.call_args
assert args[0] == Path("/mock/repo")
assert args[1] == "123"
# Verify the content contains citations
update_content = args[2]
assert "# Feature Implementation Plan" in update_content
assert "## Citations" in update_content
assert "https://docs.python.org/3/library/json.html" in update_content
assert "https://github.com/user/repo/issues/123" in update_content
# Should NOT have metrics in body
assert "## Completion Metrics" not in update_content
# Integration tests for new search grounding flow
@pytest.mark.asyncio
async def test_process_workplan_async_with_new_search_grounding(
mock_request_context, mock_genai_client
):
"""Test search grounding integration in workplan generation."""
from yellhorn_mcp.server import _get_gemini_search_tools
# Test the search tools function directly
search_tools = _get_gemini_search_tools("gemini-2.5-pro")
assert search_tools is not None
@pytest.mark.asyncio
async def test_process_workplan_async_with_search_grounding_disabled(
mock_request_context, mock_genai_client
):
"""Test that search grounding can be disabled."""
from yellhorn_mcp.server import _get_gemini_search_tools
# Test that non-Gemini models return None
search_tools = _get_gemini_search_tools("gpt-4")
assert search_tools is None
# Test that the function returns None for unsupported models
search_tools = _get_gemini_search_tools("unknown-model")
assert search_tools is None
@pytest.mark.asyncio
async def test_async_generate_content_with_config_error_handling(mock_genai_client):
"""Test async_generate_content_with_config error handling."""
pytest.skip("Removed legacy gemini_integration logic and tests")
from yellhorn_mcp.server import async_generate_content_with_config
from yellhorn_mcp.utils.git_utils import YellhornMCPError
# Test with client missing required attributes
invalid_client = MagicMock()
del invalid_client.aio
with pytest.raises(YellhornMCPError, match="does not support aio.models.generate_content"):
await async_generate_content_with_config(invalid_client, "test-model", "test prompt")
# Test successful call without generation_config
mock_response = MagicMock()
mock_genai_client.aio.models.generate_content.return_value = mock_response
result = await async_generate_content_with_config(
mock_genai_client, "test-model", "test prompt", generation_config=None
)
assert result == mock_response
mock_genai_client.aio.models.generate_content.assert_called_once_with(
model="test-model", contents="test prompt"
)
@pytest.mark.asyncio
async def test_async_generate_content_with_config_with_generation_config(mock_genai_client):
"""Test async_generate_content_with_config with generation_config parameter."""
pytest.skip("Removed legacy gemini_integration logic and tests")
from yellhorn_mcp.server import async_generate_content_with_config
# Test successful call with generation_config
mock_response = MagicMock()
mock_genai_client.aio.models.generate_content.return_value = mock_response
mock_generation_config = MagicMock()
result = await async_generate_content_with_config(
mock_genai_client, "test-model", "test prompt", generation_config=mock_generation_config
)
assert result == mock_response
mock_genai_client.aio.models.generate_content.assert_called_once_with(
model="test-model", contents="test prompt", config=mock_generation_config
)
@pytest.mark.asyncio
async def test_process_workplan_async_search_grounding_enabled(
mock_request_context, mock_genai_client
):
"""Test process_workplan_async with search grounding enabled and verify generation_config is passed."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["gemini_client"] = mock_genai_client
mock_request_context.request_context.lifespan_context["openai_client"] = None
mock_request_context.request_context.lifespan_context["use_search_grounding"] = True
with (
patch(
"yellhorn_mcp.formatters.codebase_snapshot.get_codebase_snapshot",
new_callable=AsyncMock,
) as mock_snapshot,
patch(
"yellhorn_mcp.processors.workplan_processor.format_codebase_for_prompt",
new_callable=AsyncMock,
) as mock_format,
patch(
"yellhorn_mcp.processors.workplan_processor.update_issue_with_workplan",
new_callable=AsyncMock,
) as mock_update,
patch(
"yellhorn_mcp.processors.workplan_processor.format_metrics_section"
) as mock_format_metrics,
patch(
"yellhorn_mcp.processors.workplan_processor.add_issue_comment", new_callable=AsyncMock
) as mock_add_comment,
):
mock_snapshot.return_value = (["file1.py"], {"file1.py": "content"})
mock_format.return_value = "Formatted codebase"
mock_format_metrics.return_value = (
"\n\n---\n## Completion Metrics\n* **Model Used**: `gemini-2.5-pro`"
)
# Mocks are now handled by LLM Manager
# Create mock LLM Manager
from yellhorn_mcp.llm import LLMManager
from yellhorn_mcp.llm.usage import UsageMetadata
mock_llm_manager = MagicMock(spec=LLMManager)
# Mock call_llm_with_citations to return content and usage with grounding metadata
async def mock_call_with_citations(**kwargs):
usage = UsageMetadata(
{"prompt_tokens": 1000, "completion_tokens": 500, "total_tokens": 1500}
)
return {
"content": "Generated workplan content with citations",
"usage_metadata": usage,
"grounding_metadata": MagicMock(),
"reasoning_effort": None,
}
mock_llm_manager.call_llm_with_usage = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager.call_llm_with_citations = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager._is_openai_model = MagicMock(return_value=False)
mock_llm_manager._is_gemini_model = MagicMock(return_value=True)
await process_workplan_async(
Path("/mock/repo"),
mock_llm_manager,
"gemini-2.5-pro",
"Feature Implementation Plan",
"123",
"full", # codebase_reasoning
"Create a new feature to support X", # detailed_description
ctx=mock_request_context,
github_command_func=AsyncMock(return_value="https://github.com/owner/repo/issues/123"),
git_command_func=AsyncMock(return_value="file1.py\nfile2.py"),
)
# Verify that LLM manager was called for generation with citations
mock_llm_manager.call_llm_with_citations.assert_called_once()
# Verify the final content includes citations
mock_update.assert_called_once()
update_args = mock_update.call_args
update_content = update_args[0][2]
assert "Generated workplan content with citations" in update_content
@pytest.mark.asyncio
async def test_process_judgement_async_search_grounding_enabled(
mock_request_context, mock_genai_client
):
"""Test process_judgement_async with search grounding enabled and verify generation_config is passed."""
pytest.skip("Removed legacy gemini_integration logic and tests")
from yellhorn_mcp.server import process_judgement_async
# Mock the response with grounding metadata
mock_response = MagicMock()
mock_response.text = "## Judgement Summary\nImplementation looks good."
mock_response.usage_metadata = MagicMock()
mock_response.usage_metadata.prompt_token_count = 1000
mock_response.usage_metadata.candidates_token_count = 500
mock_response.usage_metadata.total_token_count = 1500
mock_response.grounding_metadata = MagicMock()
# Mock candidates to avoid finish_reason error
mock_candidate = MagicMock()
mock_candidate.finish_reason = MagicMock()
mock_candidate.finish_reason.name = "STOP"
mock_candidate.safety_ratings = []
mock_response.candidates = [mock_candidate]
with (
patch(
"yellhorn_mcp.integrations.gemini_integration.async_generate_content_with_config"
) as mock_generate,
patch("yellhorn_mcp.utils.git_utils.update_github_issue") as mock_update_issue,
patch(
"yellhorn_mcp.processors.judgement_processor.add_issue_comment", new_callable=AsyncMock
) as mock_add_comment,
patch(
"yellhorn_mcp.processors.judgement_processor.create_judgement_subissue",
new_callable=AsyncMock,
) as mock_create_subissue,
patch(
"yellhorn_mcp.formatters.codebase_snapshot.get_codebase_snapshot",
new_callable=AsyncMock,
) as mock_snapshot,
patch(
"yellhorn_mcp.formatters.format_codebase_for_prompt",
new_callable=AsyncMock,
) as mock_format,
patch(
"yellhorn_mcp.integrations.gemini_integration._get_gemini_search_tools"
) as mock_get_tools,
patch("yellhorn_mcp.integrations.gemini_integration.add_citations") as mock_add_citations,
patch("yellhorn_mcp.processors.judgement_processor.run_git_command") as mock_run_git,
):
mock_generate.return_value = mock_response
# Mock getting the remote URL
mock_run_git.return_value = "https://github.com/user/repo"
# Mock codebase snapshot
mock_snapshot.return_value = (["file1.py"], {"file1.py": "content"})
mock_format.return_value = "Formatted codebase"
# Mock create_subissue
mock_create_subissue.return_value = "https://github.com/user/repo/issues/789"
# Mock search tools
mock_search_tools = [MagicMock()]
mock_get_tools.return_value = mock_search_tools
# Mock add_citations processing
mock_add_citations.return_value = (
"## Judgement Summary\nImplementation looks good.\n\n## Citations\n[1] Example citation"
)
# Set context for search grounding enabled
mock_request_context.request_context.lifespan_context["use_search_grounding"] = True
# Create mock LLM Manager
from yellhorn_mcp.llm import LLMManager
from yellhorn_mcp.llm.usage import UsageMetadata
mock_llm_manager = MagicMock(spec=LLMManager)
# Mock call_llm_with_citations to return content and usage with grounding metadata
async def mock_call_with_citations(**kwargs):
usage = UsageMetadata(
{"prompt_tokens": 1000, "completion_tokens": 500, "total_tokens": 1500}
)
return {
"content": "## Judgement Summary\nImplementation looks good.\n\n## Citations\n[1] Example citation",
"usage_metadata": usage,
"grounding_metadata": MagicMock(),
"reasoning_effort": None,
}
mock_llm_manager.call_llm_with_usage = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager.call_llm_with_citations = AsyncMock(side_effect=mock_call_with_citations)
mock_llm_manager._is_openai_model = MagicMock(return_value=False)
mock_llm_manager._is_gemini_model = MagicMock(return_value=True)
from datetime import datetime, timezone
await process_judgement_async(
repo_path=Path("/test/repo"),
llm_manager=mock_llm_manager,
model="gemini-2.5-pro",
workplan_content="# Workplan\n1. Do something",
diff_content="diff --git a/file.py b/file.py\n+def test(): pass",
base_ref="main",
head_ref="HEAD",
base_commit_hash="abc1234",
head_commit_hash="def5678",
parent_workplan_issue_number="123",
subissue_to_update="789",
debug=False,
codebase_reasoning="full",
ctx=mock_request_context,
)
# Verify core LLM functionality - judgement was processed with search grounding
mock_llm_manager.call_llm_with_citations.assert_called_once()
# Note: GitHub integration calls are complex to test due to dependencies
# Core judgement functionality with search grounding is verified by LLM call above
@pytest.mark.asyncio
async def test_revise_workplan(mock_request_context, mock_genai_client):
"""Test revising an existing workplan."""
# Set the mock client in the context
mock_request_context.request_context.lifespan_context["gemini_client"] = mock_genai_client
with patch("yellhorn_mcp.server.get_issue_body", new_callable=AsyncMock) as mock_get_issue:
mock_get_issue.return_value = "# Original Workplan\n## Summary\nOriginal content"
with patch(
"yellhorn_mcp.server.add_issue_comment", new_callable=AsyncMock
) as mock_add_comment:
with patch(
"yellhorn_mcp.server.run_github_command", new_callable=AsyncMock
) as mock_run_git:
# Mock getting issue URL
mock_run_git.return_value = json.dumps(
{"url": "https://github.com/user/repo/issues/123"}
)
with patch("asyncio.create_task") as mock_create_task:
# Mock the return value of create_task to avoid actual async processing
mock_task = MagicMock()
mock_create_task.return_value = mock_task
# Test revising a workplan
response = await revise_workplan(
ctx=mock_request_context,
issue_number="123",
revision_instructions="Add more detail about testing",
codebase_reasoning="full",
)
# Parse response as JSON and check contents
result = json.loads(response)
assert result["issue_url"] == "https://github.com/user/repo/issues/123"
assert result["issue_number"] == "123"
# Verify get_issue_body was called to fetch original workplan
mock_get_issue.assert_called_once_with(Path("/mock/repo"), "123")
# Verify submission comment was added
mock_add_comment.assert_called_once()
comment_args = mock_add_comment.call_args
assert comment_args[0][0] == Path("/mock/repo") # repo_path
assert comment_args[0][1] == "123" # issue_number
submission_comment = comment_args[0][2] # comment content
# Verify the submission comment contains expected metadata
assert "## 🚀 Revising workplan..." in submission_comment
assert "**Model**: `gemini-2.5-pro`" in submission_comment
assert "**Codebase Reasoning**: `full`" in submission_comment
# Check that the process_revision_async task is created
args, kwargs = mock_create_task.call_args
coroutine = args[0]
assert coroutine.__name__ == "process_revision_async"
# Close the coroutine to prevent RuntimeWarning
coroutine.close()
# Reset mocks for next test
mock_get_issue.reset_mock()
mock_add_comment.reset_mock()
mock_run_git.reset_mock()
mock_create_task.reset_mock()
# Test with non-existent issue
mock_get_issue.return_value = None
with pytest.raises(Exception) as exc_info:
await revise_workplan(
ctx=mock_request_context,
issue_number="999",
revision_instructions="Update something",
)
assert "Could not retrieve workplan for issue #999" in str(exc_info.value)
# Additional tests for better coverage
@pytest.mark.asyncio
async def test_app_lifespan_gemini_model():
"""Test app_lifespan with Gemini model configuration."""
from mcp.server.fastmcp import FastMCP
from yellhorn_mcp.server import app_lifespan
mock_server = MagicMock(spec=FastMCP)
with (
patch("os.getenv") as mock_getenv,
patch("pathlib.Path.resolve") as mock_resolve,
patch("yellhorn_mcp.server.is_git_repository", return_value=True),
patch("yellhorn_mcp.server.genai.Client") as mock_gemini_client,
):
# Mock environment variables for Gemini model
def getenv_side_effect(key, default=None):
env_vars = {
"REPO_PATH": "/test/repo",
"YELLHORN_MCP_MODEL": "gemini-2.5-pro",
"YELLHORN_MCP_SEARCH": "on",
"GEMINI_API_KEY": "test-gemini-key",
}
return env_vars.get(key, default)
mock_getenv.side_effect = getenv_side_effect
mock_resolve.return_value = Path("/test/repo")
async with app_lifespan(mock_server) as lifespan_context:
assert lifespan_context["repo_path"] == Path("/test/repo")
assert lifespan_context["model"] == "gemini-2.5-pro"
assert lifespan_context["use_search_grounding"] is True
assert lifespan_context["gemini_client"] is not None
assert lifespan_context["openai_client"] is None
assert lifespan_context["llm_manager"] is not None
@pytest.mark.asyncio
async def test_app_lifespan_openai_model():
"""Test app_lifespan with OpenAI model configuration."""
from mcp.server.fastmcp import FastMCP
from yellhorn_mcp.server import app_lifespan
mock_server = MagicMock(spec=FastMCP)
with (
patch("os.getenv") as mock_getenv,
patch("pathlib.Path.resolve") as mock_resolve,
patch("yellhorn_mcp.server.is_git_repository", return_value=True),
patch("httpx.AsyncClient") as mock_httpx,
patch("yellhorn_mcp.server.AsyncOpenAI") as mock_openai_client,
):
# Mock environment variables for OpenAI model
def getenv_side_effect(key, default=None):
env_vars = {
"REPO_PATH": "/test/repo",
"YELLHORN_MCP_MODEL": "gpt-4o",
"OPENAI_API_KEY": "test-openai-key",
}
return env_vars.get(key, default)
mock_getenv.side_effect = getenv_side_effect
mock_resolve.return_value = Path("/test/repo")
async with app_lifespan(mock_server) as lifespan_context:
assert lifespan_context["repo_path"] == Path("/test/repo")
assert lifespan_context["model"] == "gpt-4o"
assert lifespan_context["use_search_grounding"] is False # Disabled for OpenAI
assert lifespan_context["gemini_client"] is None
assert lifespan_context["openai_client"] is not None
assert lifespan_context["xai_client"] is None
assert lifespan_context["llm_manager"] is not None
@pytest.mark.asyncio
async def test_app_lifespan_grok_model(caplog):
"""Test app_lifespan initializes Grok models with xAI credentials."""
from mcp.server.fastmcp import FastMCP
from yellhorn_mcp.server import app_lifespan
mock_server = MagicMock(spec=FastMCP)
caplog.set_level("INFO")
with (
patch("os.getenv") as mock_getenv,
patch("pathlib.Path.resolve") as mock_resolve,
patch("yellhorn_mcp.server.is_git_repository", return_value=True),
patch("yellhorn_mcp.server.AsyncXAI") as mock_async_xai,
):
def getenv_side_effect(key, default=None):
env_vars = {
"REPO_PATH": "/test/repo",
"YELLHORN_MCP_MODEL": "grok-4",
"XAI_API_KEY": "test-xai-key",
"XAI_API_BASE_URL": "https://mock.x.ai/v1",
}
return env_vars.get(key, default)
mock_getenv.side_effect = getenv_side_effect
mock_resolve.return_value = Path("/test/repo")
async with app_lifespan(mock_server) as lifespan_context:
assert lifespan_context["model"] == "grok-4"
assert lifespan_context["use_search_grounding"] is False
assert lifespan_context["gemini_client"] is None
assert lifespan_context["openai_client"] is None
assert lifespan_context["xai_client"] is not None
mock_async_xai.assert_called_once_with(api_key="test-xai-key", api_host="mock.x.ai")
assert any("Initializing Grok client" in record.message for record in caplog.records)
@pytest.mark.asyncio
async def test_app_lifespan_missing_gemini_api_key():
"""Test app_lifespan raises error when Gemini API key is missing."""
from mcp.server.fastmcp import FastMCP
from yellhorn_mcp.server import app_lifespan
mock_server = MagicMock(spec=FastMCP)
with patch("os.getenv") as mock_getenv:
# Mock environment variables without Gemini API key
def getenv_side_effect(key, default=None):
env_vars = {
"REPO_PATH": "/test/repo",
"YELLHORN_MCP_MODEL": "gemini-2.5-pro",
# GEMINI_API_KEY is missing
}
return env_vars.get(key, default)
mock_getenv.side_effect = getenv_side_effect
with pytest.raises(ValueError, match="GEMINI_API_KEY is required for Gemini models"):
async with app_lifespan(mock_server) as _:
pass
@pytest.mark.asyncio
async def test_app_lifespan_missing_openai_api_key():
"""Test app_lifespan raises error when OpenAI API key is missing."""
from mcp.server.fastmcp import FastMCP
from yellhorn_mcp.server import app_lifespan
mock_server = MagicMock(spec=FastMCP)
with patch("os.getenv") as mock_getenv:
# Mock environment variables without OpenAI API key
def getenv_side_effect(key, default=None):
env_vars = {
"REPO_PATH": "/test/repo",
"YELLHORN_MCP_MODEL": "gpt-4o",
# OPENAI_API_KEY is missing
}
return env_vars.get(key, default)
mock_getenv.side_effect = getenv_side_effect
with pytest.raises(ValueError, match="OPENAI_API_KEY is required for OpenAI models"):
async with app_lifespan(mock_server) as _:
pass
@pytest.mark.asyncio
async def test_app_lifespan_grok_missing_sdk():
"""Grok models should raise if xai-sdk is unavailable."""
from mcp.server.fastmcp import FastMCP
from yellhorn_mcp.server import app_lifespan
mock_server = MagicMock(spec=FastMCP)
with (
patch("os.getenv") as mock_getenv,
patch("yellhorn_mcp.server.AsyncXAI", None),
):
def getenv_side_effect(key, default=None):
env_vars = {
"REPO_PATH": "/test/repo",
"YELLHORN_MCP_MODEL": "grok-4",
"XAI_API_KEY": "test-xai-key",
}
return env_vars.get(key, default)
mock_getenv.side_effect = getenv_side_effect
with pytest.raises(ValueError, match="xai-sdk is required for Grok models"):
async with app_lifespan(mock_server) as _:
pass
@pytest.mark.asyncio
async def test_app_lifespan_invalid_repository():
"""Test app_lifespan raises error for invalid repository path."""
from mcp.server.fastmcp import FastMCP
from yellhorn_mcp.server import app_lifespan
mock_server = MagicMock(spec=FastMCP)
with (
patch("os.getenv") as mock_getenv,
patch("pathlib.Path.resolve") as mock_resolve,
patch("yellhorn_mcp.server.is_git_repository", return_value=False),
):
# Mock environment variables
def getenv_side_effect(key, default=None):
env_vars = {
"REPO_PATH": "/not/a/git/repo",
"YELLHORN_MCP_MODEL": "gemini-2.5-pro",
"GEMINI_API_KEY": "test-key",
}
return env_vars.get(key, default)
mock_getenv.side_effect = getenv_side_effect
mock_resolve.return_value = Path("/not/a/git/repo")
with pytest.raises(ValueError, match="Path .* is not a Git repository"):
async with app_lifespan(mock_server) as _:
pass
@pytest.mark.asyncio
async def test_app_lifespan_search_grounding_disabled():
"""Test app_lifespan with search grounding explicitly disabled."""
from mcp.server.fastmcp import FastMCP
from yellhorn_mcp.server import app_lifespan
mock_server = MagicMock(spec=FastMCP)
with (
patch("os.getenv") as mock_getenv,
patch("pathlib.Path.resolve") as mock_resolve,
patch("yellhorn_mcp.server.is_git_repository", return_value=True),
patch("yellhorn_mcp.server.genai.Client") as mock_gemini_client,
):
# Mock environment variables with search grounding disabled
def getenv_side_effect(key, default=None):
env_vars = {
"REPO_PATH": "/test/repo",
"YELLHORN_MCP_MODEL": "gemini-2.5-pro",
"YELLHORN_MCP_SEARCH": "off", # Explicitly disabled
"GEMINI_API_KEY": "test-gemini-key",
}
return env_vars.get(key, default)
mock_getenv.side_effect = getenv_side_effect
mock_resolve.return_value = Path("/test/repo")
async with app_lifespan(mock_server) as lifespan_context:
assert lifespan_context["use_search_grounding"] is False
@pytest.mark.asyncio
async def test_app_lifespan_o_series_model():
"""Test app_lifespan correctly identifies 'o' series models as OpenAI."""
from mcp.server.fastmcp import FastMCP
from yellhorn_mcp.server import app_lifespan
mock_server = MagicMock(spec=FastMCP)
with (
patch("os.getenv") as mock_getenv,
patch("pathlib.Path.resolve") as mock_resolve,
patch("yellhorn_mcp.server.is_git_repository", return_value=True),
patch("httpx.AsyncClient") as mock_httpx,
patch("yellhorn_mcp.server.AsyncOpenAI") as mock_openai_client,
):
# Mock environment variables for 'o' series model
def getenv_side_effect(key, default=None):
env_vars = {
"REPO_PATH": "/test/repo",
"YELLHORN_MCP_MODEL": "o3-deep-research",
"OPENAI_API_KEY": "test-openai-key",
}
return env_vars.get(key, default)
mock_getenv.side_effect = getenv_side_effect
mock_resolve.return_value = Path("/test/repo")
async with app_lifespan(mock_server) as lifespan_context:
assert lifespan_context["model"] == "o3-deep-research"
assert lifespan_context["use_search_grounding"] is False # Disabled for OpenAI
assert lifespan_context["openai_client"] is not None
assert lifespan_context["gemini_client"] is None
@pytest.mark.asyncio
async def test_create_workplan_no_llm_manager():
"""Test create_workplan handles missing LLM manager gracefully."""
from yellhorn_mcp.server import create_workplan
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
"llm_manager": None, # No LLM manager
"model": "gpt-4o",
"use_search_grounding": False,
}
mock_ctx.log = AsyncMock()
with (
patch("yellhorn_mcp.server.create_github_issue") as mock_create_issue,
patch("yellhorn_mcp.server.add_issue_comment") as mock_add_comment,
):
mock_create_issue.return_value = {
"number": "123",
"url": "https://github.com/user/repo/issues/123",
}
result = await create_workplan(
ctx=mock_ctx,
title="Test Workplan",
detailed_description="Test description",
)
result_data = json.loads(result)
assert result_data["issue_number"] == "123"
# Should add submission comment (LLM manager error happens in background task)
mock_add_comment.assert_called_once()
@pytest.mark.asyncio
async def test_create_workplan_with_urls():
"""Test create_workplan with URLs in description."""
from yellhorn_mcp.server import create_workplan
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
"llm_manager": MagicMock(),
"model": "gpt-4o",
"use_search_grounding": False,
}
mock_ctx.log = AsyncMock()
with (
patch("yellhorn_mcp.server.create_github_issue") as mock_create_issue,
patch("yellhorn_mcp.server.extract_urls") as mock_extract_urls,
patch("yellhorn_mcp.server.format_submission_comment") as mock_format_comment,
patch("yellhorn_mcp.server.add_issue_comment") as mock_add_comment,
patch("asyncio.create_task") as mock_create_task,
):
mock_create_issue.return_value = {
"number": "123",
"url": "https://github.com/user/repo/issues/123",
}
mock_extract_urls.return_value = ["https://example.com", "https://github.com/user/repo"]
mock_format_comment.return_value = "Submission comment"
description_with_urls = (
"Test description with https://example.com and https://github.com/user/repo"
)
result = await create_workplan(
ctx=mock_ctx,
title="Test Workplan",
detailed_description=description_with_urls,
)
result_data = json.loads(result)
assert result_data["issue_number"] == "123"
mock_extract_urls.assert_called_once_with(description_with_urls)
# Should create async task for processing
mock_create_task.assert_called_once()
@pytest.mark.asyncio
async def test_get_workplan_success():
"""Test get_workplan successfully retrieves workplan content."""
from yellhorn_mcp.server import get_workplan
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
}
with patch("yellhorn_mcp.server.get_issue_body") as mock_get_body:
mock_get_body.return_value = "# Test Workplan\n\nThis is a test workplan."
result = await get_workplan(ctx=mock_ctx, issue_number="123")
assert result == "# Test Workplan\n\nThis is a test workplan."
mock_get_body.assert_called_once_with(Path("/test/repo"), "123")
@pytest.mark.asyncio
async def test_get_workplan_failure():
"""Test get_workplan handles errors when retrieving workplan."""
from yellhorn_mcp.server import get_workplan
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
}
with patch("yellhorn_mcp.server.get_issue_body") as mock_get_body:
mock_get_body.side_effect = Exception("GitHub API error")
with pytest.raises(Exception, match="GitHub API error"):
await get_workplan(ctx=mock_ctx, issue_number="123")
@pytest.mark.asyncio
async def test_curate_context_success():
"""Test curate_context successfully processes context curation."""
from yellhorn_mcp.server import curate_context
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
"llm_manager": MagicMock(),
"model": "gemini-2.5-pro",
}
mock_ctx.log = AsyncMock()
with patch("yellhorn_mcp.server.process_context_curation_async") as mock_process:
mock_process.return_value = None
result = await curate_context(
ctx=mock_ctx,
user_task="Implement authentication system",
codebase_reasoning="file_structure",
)
result_data = json.loads(result)
assert result_data["status"] == "✅ Context curation completed successfully"
mock_process.assert_called_once()
@pytest.mark.asyncio
async def test_curate_context_no_llm_manager():
"""Test curate_context handles missing LLM manager."""
from yellhorn_mcp.server import curate_context
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
"llm_manager": None, # No LLM manager
"model": "gemini-2.5-pro",
}
mock_ctx.log = AsyncMock()
with pytest.raises(YellhornMCPError, match="LLM Manager not initialized"):
await curate_context(
ctx=mock_ctx,
user_task="Implement authentication system",
codebase_reasoning="file_structure",
)
@pytest.mark.asyncio
async def test_curate_context_with_optional_params():
"""Test curate_context with optional parameters."""
from yellhorn_mcp.server import curate_context
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
"llm_manager": MagicMock(),
"model": "gemini-2.5-pro",
}
mock_ctx.log = AsyncMock()
with patch("yellhorn_mcp.server.process_context_curation_async") as mock_process:
mock_process.return_value = None
result = await curate_context(
ctx=mock_ctx,
user_task="Implement authentication system",
codebase_reasoning="lsp",
ignore_file_path=".myignore",
output_path=".mycontext",
)
result_data = json.loads(result)
assert result_data["status"] == "✅ Context curation completed successfully"
# Check that optional parameters were passed
call_args = mock_process.call_args
assert call_args.kwargs["codebase_reasoning"] == "lsp"
assert call_args.kwargs["output_path"] == ".mycontext"
@pytest.mark.asyncio
async def test_judge_workplan_basic():
"""Test judge_workplan basic functionality."""
from yellhorn_mcp.server import judge_workplan
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
"llm_manager": MagicMock(),
"model": "gemini-2.5-pro",
"use_search_grounding": True,
}
mock_ctx.log = AsyncMock()
with (
patch("yellhorn_mcp.server.get_issue_body") as mock_get_body,
patch("yellhorn_mcp.server.get_git_diff") as mock_get_diff,
patch("asyncio.create_task") as mock_create_task,
patch("yellhorn_mcp.server.run_git_command") as mock_git,
):
mock_get_body.return_value = "# Workplan content"
mock_get_diff.return_value = "diff content"
mock_git.return_value = "abc123" # commit hash
result = await judge_workplan(
ctx=mock_ctx,
issue_number="123",
base_ref="main",
head_ref="feature",
)
result_data = json.loads(result)
assert "subissue_url" in result_data
assert "subissue_number" in result_data
mock_create_task.assert_called_once()
@pytest.mark.asyncio
async def test_judge_workplan_missing_workplan():
"""Test judge_workplan handles missing workplan gracefully."""
from yellhorn_mcp.server import judge_workplan
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
"llm_manager": MagicMock(),
"model": "gemini-2.5-pro",
"use_search_grounding": True,
}
mock_ctx.log = AsyncMock()
with patch("yellhorn_mcp.server.get_issue_body") as mock_get_body:
mock_get_body.side_effect = Exception("Issue not found")
with pytest.raises(Exception, match="Issue not found"):
await judge_workplan(
ctx=mock_ctx,
issue_number="999",
base_ref="main",
head_ref="feature",
)
@pytest.mark.asyncio
async def test_judge_workplan_no_llm_manager():
"""Test judge_workplan with no LLM manager."""
from yellhorn_mcp.server import judge_workplan
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
"llm_manager": None, # No LLM manager
"model": "gemini-2.5-pro",
"use_search_grounding": True,
}
mock_ctx.log = AsyncMock()
with (
patch("yellhorn_mcp.server.get_issue_body") as mock_get_body,
patch("yellhorn_mcp.server.get_git_diff") as mock_get_diff,
patch("yellhorn_mcp.server.run_git_command") as mock_git,
):
mock_get_body.return_value = "# Workplan content"
mock_get_diff.return_value = "diff content"
mock_git.return_value = "abc123"
result = await judge_workplan(
ctx=mock_ctx,
issue_number="123",
base_ref="main",
head_ref="feature",
)
# Should still return JSON with subissue info even without LLM manager
result_data = json.loads(result)
assert "subissue_url" in result_data
assert "subissue_number" in result_data
@pytest.mark.asyncio
async def test_judge_workplan_with_subissue():
"""Test judge_workplan with existing subissue to update."""
from yellhorn_mcp.server import judge_workplan
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
"llm_manager": MagicMock(),
"model": "gemini-2.5-pro",
"use_search_grounding": True,
}
mock_ctx.log = AsyncMock()
with (
patch("yellhorn_mcp.server.get_issue_body") as mock_get_body,
patch("yellhorn_mcp.server.get_git_diff") as mock_get_diff,
patch("asyncio.create_task") as mock_create_task,
patch("yellhorn_mcp.server.run_git_command") as mock_git,
):
mock_get_body.return_value = "# Workplan content"
mock_get_diff.return_value = "diff content"
mock_git.return_value = "abc123"
result = await judge_workplan(
ctx=mock_ctx,
issue_number="123",
base_ref="main",
head_ref="feature",
subissue_to_update="456", # Existing subissue
)
result_data = json.loads(result)
assert "subissue_url" in result_data
assert "subissue_number" in result_data
mock_create_task.assert_called_once()
@pytest.mark.asyncio
async def test_judge_workplan_with_pr_url():
"""Test judge_workplan with PR URL instead of refs."""
from yellhorn_mcp.server import judge_workplan
mock_ctx = MagicMock()
mock_ctx.request_context.lifespan_context = {
"repo_path": Path("/test/repo"),
"llm_manager": MagicMock(),
"model": "gemini-2.5-pro",
"use_search_grounding": True,
}
mock_ctx.log = AsyncMock()
with (
patch("yellhorn_mcp.server.get_issue_body") as mock_get_body,
patch("yellhorn_mcp.server.get_github_pr_diff") as mock_get_pr_diff,
patch("asyncio.create_task") as mock_create_task,
):
mock_get_body.return_value = "# Workplan content"
mock_get_pr_diff.return_value = "PR diff content"
result = await judge_workplan(
ctx=mock_ctx,
issue_number="123",
pr_url="https://github.com/user/repo/pull/456",
)
result_data = json.loads(result)
assert "subissue_url" in result_data
assert "subissue_number" in result_data
mock_get_pr_diff.assert_called_once_with(
Path("/test/repo"), "https://github.com/user/repo/pull/456"
)
mock_create_task.assert_called_once()
