diff --git a/pyproject.toml b/pyproject.toml index 926332dd0..e48023ef2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,6 @@ dependencies = [ "codeowners>=0.6.0", "unidiff>=0.7.5", "datamodel-code-generator>=0.26.5", - "mcp[cli]==1.9.4", "fastmcp>=2.9.0", # Utility dependencies "colorlog>=6.9.0", diff --git a/src/codegen/cli/commands/mcp/__init__.py b/src/codegen/cli/commands/mcp/__init__.py index 139597f9c..e69de29bb 100644 --- a/src/codegen/cli/commands/mcp/__init__.py +++ b/src/codegen/cli/commands/mcp/__init__.py @@ -1,2 +0,0 @@ - - diff --git a/src/codegen/cli/mcp/resources/system_prompt.py b/src/codegen/cli/mcp/resources/system_prompt.py index 9c7e23c6b..76f4eccaf 100644 --- a/src/codegen/cli/mcp/resources/system_prompt.py +++ b/src/codegen/cli/mcp/resources/system_prompt.py @@ -8439,10 +8439,10 @@ class FeatureFlags: ```python python # Import datetime for timestamp -from datetime import datetime +from datetime import datetime, timezone # Get current timestamp -timestamp = datetime.now().strftime("%B %d, %Y") +timestamp = datetime.now(timezone.utc).strftime("%B %d, %Y") print("📚 Generating and Updating Function Documentation") diff --git a/src/codegen/git/repo_operator/repo_operator.py b/src/codegen/git/repo_operator/repo_operator.py index f3bf2776f..ff81e62c9 100644 --- a/src/codegen/git/repo_operator/repo_operator.py +++ b/src/codegen/git/repo_operator/repo_operator.py @@ -566,7 +566,7 @@ def emptydir(self, path: str) -> None: if os.path.isfile(file_path): os.remove(file_path) - def get_file(self, path: str) -> str: + def get_file(self, path: str) -> str | None: """Returns the contents of a file""" file_path = self.abspath(path) try: @@ -621,7 +621,7 @@ def get_filepaths_for_repo(self, ignore_list): decoded_filepath, _ = codecs.escape_decode(raw_filepath) # Step 4: Decode those bytes as UTF-8 to get the actual Unicode text - filepath = decoded_filepath.decode("utf-8") + filepath = decoded_filepath.decode("utf-8") # type: ignore[union-attr] # Step 5: Replace the original filepath with the decoded filepath filepaths[i] = filepath @@ -754,7 +754,7 @@ def get_pr_data(self, pr_number: int) -> dict: """Returns the data associated with a PR""" return self.remote_git_repo.get_pr_data(pr_number) - def create_pr_comment(self, pr_number: int, body: str) -> IssueComment: + def create_pr_comment(self, pr_number: int, body: str) -> IssueComment | None: """Create a general comment on a pull request. Args: @@ -765,6 +765,7 @@ def create_pr_comment(self, pr_number: int, body: str) -> IssueComment: if pr: comment = self.remote_git_repo.create_issue_comment(pr, body) return comment + return None def create_pr_review_comment( self, diff --git a/src/codegen/git/schemas/repo_config.py b/src/codegen/git/schemas/repo_config.py index f94e85592..841db99cc 100644 --- a/src/codegen/git/schemas/repo_config.py +++ b/src/codegen/git/schemas/repo_config.py @@ -31,11 +31,13 @@ class RepoConfig(BaseModel): @classmethod def from_envs(cls) -> "RepoConfig": default_repo_config = RepositoryConfig() + path = default_repo_config.path or os.getcwd() + language = default_repo_config.language or "python" return RepoConfig( name=default_repo_config.name, full_name=default_repo_config.full_name, - base_dir=os.path.dirname(default_repo_config.path), - language=ProgrammingLanguage(default_repo_config.language.upper()), + base_dir=os.path.dirname(path), + language=ProgrammingLanguage(language.upper()), ) @classmethod diff --git a/src/codegen/git/utils/language.py b/src/codegen/git/utils/language.py index 551ac4212..19a6e5647 100644 --- a/src/codegen/git/utils/language.py +++ b/src/codegen/git/utils/language.py @@ -46,13 +46,20 @@ def _determine_language_by_file_count(folder_path: str) -> ProgrammingLanguage: ProgrammingLanguage: The dominant programming language, or OTHER if no matching files found or if less than MIN_LANGUAGE_RATIO of files match the dominant language """ - from codegen.sdk.python import PyFile - from codegen.sdk.typescript.file import TSFile - - EXTENSIONS = { - ProgrammingLanguage.PYTHON: PyFile.get_extensions(), - ProgrammingLanguage.TYPESCRIPT: TSFile.get_extensions(), - } + try: + from codegen.sdk.python import PyFile + from codegen.sdk.typescript.file import TSFile + + EXTENSIONS = { + ProgrammingLanguage.PYTHON: PyFile.get_extensions(), + ProgrammingLanguage.TYPESCRIPT: TSFile.get_extensions(), + } + except ImportError: + # Fallback to hardcoded extensions if SDK modules are not available + EXTENSIONS = { + ProgrammingLanguage.PYTHON: [".py", ".pyx", ".pyi"], + ProgrammingLanguage.TYPESCRIPT: [".ts", ".tsx", ".js", ".jsx"], + } folder = Path(folder_path) if not folder.exists() or not folder.is_dir(): diff --git a/src/codegen/git/utils/pr_review.py b/src/codegen/git/utils/pr_review.py index ffb3f52f0..03e5b23b4 100644 --- a/src/codegen/git/utils/pr_review.py +++ b/src/codegen/git/utils/pr_review.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING -from github import Repository from github.PullRequest import PullRequest +from github.Repository import Repository from unidiff import PatchSet from codegen.git.models.pull_request_context import PullRequestContext @@ -97,7 +97,7 @@ class CodegenPR: _op: RepoOperator # =====[ Computed ]===== - _modified_file_ranges: dict[str, list[tuple[int, int]]] = None + _modified_file_ranges: dict[str, list[tuple[int, int]]] | None = None def __init__(self, op: RepoOperator, codebase: "Codebase", pr: PullRequest): self._op = op diff --git a/src/codegen/shared/decorators/docs.py b/src/codegen/shared/decorators/docs.py index a7a1aaa49..a72add8d1 100644 --- a/src/codegen/shared/decorators/docs.py +++ b/src/codegen/shared/decorators/docs.py @@ -2,14 +2,14 @@ import inspect from collections.abc import Callable from dataclasses import dataclass -from typing import TypeVar +from typing import Any, TypeVar @dataclass class DocumentedObject: name: str module: str - object: any + object: Any def __lt__(self, other): return self.module < other.module diff --git a/tests/cli/mcp/__init__.py b/tests/cli/mcp/__init__.py index 8b1378917..e69de29bb 100644 --- a/tests/cli/mcp/__init__.py +++ b/tests/cli/mcp/__init__.py @@ -1 +0,0 @@ - diff --git a/tests/cli/mcp/test_mcp_configuration.py b/tests/cli/mcp/test_mcp_configuration.py new file mode 100644 index 000000000..e0357d7bb --- /dev/null +++ b/tests/cli/mcp/test_mcp_configuration.py @@ -0,0 +1,376 @@ +"""Tests for MCP configuration scenarios.""" + +import os +import subprocess +import time +from pathlib import Path +from unittest.mock import patch + + +class TestMCPConfiguration: + """Test MCP configuration scenarios.""" + + def test_mcp_server_default_configuration(self): + """Test MCP server with default configuration.""" + codegen_path = Path(__file__).parent.parent.parent.parent / "src" + venv_python = Path(__file__).parent.parent.parent.parent / ".venv" / "bin" / "python" + + env = os.environ.copy() + env["PYTHONPATH"] = str(codegen_path) + + process = subprocess.Popen( + [str(venv_python), "-c", "from codegen.cli.cli import main; main(['mcp'])"], + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + try: + # Give it time to start + time.sleep(2) + + # Server should start with default configuration + assert process.poll() is None, "Server should start with default configuration" + + finally: + # Cleanup + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + def test_mcp_server_stdio_transport_explicit(self): + """Test MCP server with explicit stdio transport.""" + codegen_path = Path(__file__).parent.parent.parent.parent / "src" + venv_python = Path(__file__).parent.parent.parent.parent / ".venv" / "bin" / "python" + + env = os.environ.copy() + env["PYTHONPATH"] = str(codegen_path) + + process = subprocess.Popen( + [str(venv_python), "-c", "from codegen.cli.cli import main; main(['mcp', '--transport', 'stdio'])"], + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + try: + # Give it time to start + time.sleep(2) + + # Server should start with stdio transport + assert process.poll() is None, "Server should start with stdio transport" + + finally: + # Cleanup + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + def test_mcp_server_http_transport_configuration(self): + """Test MCP server with HTTP transport configuration.""" + codegen_path = Path(__file__).parent.parent.parent.parent / "src" + venv_python = Path(__file__).parent.parent.parent.parent / ".venv" / "bin" / "python" + + env = os.environ.copy() + env["PYTHONPATH"] = str(codegen_path) + + # Test HTTP transport (should fall back to stdio for now) + process = subprocess.Popen( + [str(venv_python), "-c", "from codegen.cli.cli import main; main(['mcp', '--transport', 'http', '--host', '127.0.0.1', '--port', '8080'])"], + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + try: + # Give it time to start + time.sleep(2) + + # Server should start (even if it falls back to stdio) + assert process.poll() is None, "Server should start with HTTP transport configuration" + + finally: + # Cleanup + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + def test_mcp_server_custom_host_port(self): + """Test MCP server with custom host and port.""" + codegen_path = Path(__file__).parent.parent.parent.parent / "src" + venv_python = Path(__file__).parent.parent.parent.parent / ".venv" / "bin" / "python" + + env = os.environ.copy() + env["PYTHONPATH"] = str(codegen_path) + + process = subprocess.Popen( + [str(venv_python), "-c", "from codegen.cli.cli import main; main(['mcp', '--host', '0.0.0.0', '--port', '9000'])"], + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + try: + # Give it time to start + time.sleep(2) + + # Server should start with custom host and port + assert process.poll() is None, "Server should start with custom host and port" + + finally: + # Cleanup + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + def test_mcp_server_environment_variables(self): + """Test MCP server with various environment variables.""" + codegen_path = Path(__file__).parent.parent.parent.parent / "src" + venv_python = Path(__file__).parent.parent.parent.parent / ".venv" / "bin" / "python" + + env = os.environ.copy() + env["PYTHONPATH"] = str(codegen_path) + env["CODEGEN_API_KEY"] = "test-api-key-123" + env["CODEGEN_API_BASE_URL"] = "https://custom.api.codegen.com" + + process = subprocess.Popen( + [str(venv_python), "-c", "from codegen.cli.cli import main; main(['mcp', '--transport', 'stdio'])"], + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + try: + # Give it time to start + time.sleep(2) + + # Server should start with custom environment variables + assert process.poll() is None, "Server should start with custom environment variables" + + finally: + # Cleanup + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + def test_api_client_configuration_with_env_vars(self): + """Test API client configuration with environment variables.""" + from codegen.cli.mcp.server import get_api_client + + with patch.dict(os.environ, {"CODEGEN_API_KEY": "test-key-456", "CODEGEN_API_BASE_URL": "https://test.api.codegen.com"}): + try: + api_client, agents_api, orgs_api, users_api = get_api_client() + + # Should return configured API clients + assert api_client is not None + assert agents_api is not None + assert orgs_api is not None + assert users_api is not None + + # Check that configuration was applied + assert api_client.configuration.host == "https://test.api.codegen.com" + + except Exception as e: + # If API client is not available, that's expected in test environment + if "codegen-api-client is not available" not in str(e): + raise + + def test_api_client_configuration_defaults(self): + """Test API client configuration with default values.""" + from codegen.cli.mcp.server import get_api_client + + # Clear environment variables to test defaults + with patch.dict(os.environ, {}, clear=True): + try: + api_client, agents_api, orgs_api, users_api = get_api_client() + + # Should use default base URL + assert api_client.configuration.host == "https://api.codegen.com" + + except Exception as e: + # If API client is not available, that's expected in test environment + if "codegen-api-client is not available" not in str(e): + raise + + def test_mcp_server_configuration_validation(self): + """Test MCP server configuration validation.""" + # Test that the function has the expected parameters + import inspect + + from codegen.cli.commands.mcp.main import mcp + + sig = inspect.signature(mcp) + + # Check that all expected parameters are present + expected_params = ["host", "port", "transport"] + for param in expected_params: + assert param in sig.parameters, f"Parameter {param} not found in mcp function signature" + + # Check parameter defaults + assert sig.parameters["host"].default == "localhost" + assert sig.parameters["port"].default is None + assert sig.parameters["transport"].default == "stdio" + + def test_transport_validation_in_command(self): + """Test transport validation in the MCP command.""" + from codegen.cli.commands.mcp.main import mcp + + # This test would ideally call the function with invalid transport + # but since it would try to actually run the server, we'll test the validation logic + # by checking that the function exists and has the right structure + + # The function should exist and be callable + assert callable(mcp) + + # The validation logic is in the function body, so we can't easily test it + # without actually running the server, which we do in integration tests + + def test_server_configuration_object_creation(self): + """Test server configuration object creation.""" + from codegen.cli.mcp.server import mcp + + # Check that the FastMCP server was created with correct configuration + assert mcp.name == "codegen-mcp" + assert mcp.instructions is not None + assert "MCP server for the Codegen platform" in mcp.instructions + + # Check that tools and resources are registered + assert len(mcp._tool_manager._tools) > 0, "Server should have tools registered" + assert len(mcp._resource_manager._resources) > 0, "Server should have resources registered" + + def test_server_instructions_configuration(self): + """Test server instructions configuration.""" + from codegen.cli.mcp.server import mcp + + instructions = mcp.instructions + assert instructions is not None + # Should contain key information about the server's purpose + assert "MCP server" in instructions + assert "Codegen" in instructions + assert "tools" in instructions + assert "resources" in instructions + + def test_global_api_client_singleton_behavior(self): + """Test global API client singleton behavior.""" + # Reset global state + import codegen.cli.mcp.server + from codegen.cli.mcp.server import get_api_client + + codegen.cli.mcp.server._api_client = None + codegen.cli.mcp.server._agents_api = None + codegen.cli.mcp.server._organizations_api = None + codegen.cli.mcp.server._users_api = None + + try: + # First call should create the client + client1 = get_api_client() + + # Second call should return the same client + client2 = get_api_client() + + # Should be the same objects (singleton behavior) + assert client1[0] is client2[0], "API client should be singleton" + assert client1[1] is client2[1], "Agents API should be singleton" + assert client1[2] is client2[2], "Organizations API should be singleton" + assert client1[3] is client2[3], "Users API should be singleton" + + except Exception as e: + # If API client is not available, that's expected in test environment + if "codegen-api-client is not available" not in str(e): + raise + + def test_conditional_tool_registration(self): + """Test conditional tool registration based on available imports.""" + from codegen.cli.mcp.server import mcp, LEGACY_IMPORTS_AVAILABLE # type: ignore[attr-defined] + tool_names = list(mcp._tool_manager._tools.keys()) + + if LEGACY_IMPORTS_AVAILABLE: + # Legacy tools should be available + assert "ask_codegen_sdk" in tool_names, "ask_codegen_sdk should be available when legacy imports are available" + assert "improve_codemod" in tool_names, "improve_codemod should be available when legacy imports are available" + else: + # Legacy tools should not be available + assert "ask_codegen_sdk" not in tool_names, "ask_codegen_sdk should not be available when legacy imports are unavailable" + assert "improve_codemod" not in tool_names, "improve_codemod should not be available when legacy imports are unavailable" + + # Core tools should always be available + core_tools = ["generate_codemod", "create_agent_run", "get_agent_run", "get_organizations", "get_users", "get_user"] + for tool in core_tools: + assert tool in tool_names, f"Core tool {tool} should always be available" + + def test_server_name_and_metadata(self): + """Test server name and metadata configuration.""" + from codegen.cli.mcp.server import mcp + + # Check server metadata + assert mcp.name == "codegen-mcp" + + # Check that the server has the expected structure + assert hasattr(mcp, "_tool_manager") + assert hasattr(mcp, "_resource_manager") + assert hasattr(mcp, "instructions") + + def test_resource_configuration_consistency(self): + """Test resource configuration consistency.""" + from codegen.cli.mcp.server import mcp + + # All resources should have URIs, descriptions, and MIME types + resources = mcp._resource_manager._resources + for uri, resource in resources.items(): + assert hasattr(resource, "description"), "Resource should have description" + assert hasattr(resource, "mime_type"), "Resource should have MIME type" + assert hasattr(resource, "fn"), "Resource should have function" + + # URI should be a string + assert isinstance(uri, str), "Resource URI should be string" + assert len(uri) > 0, "Resource URI should not be empty" + + # MIME type should be valid + valid_mime_types = ["text/plain", "application/json", "text/html", "application/xml"] + assert resource.mime_type in valid_mime_types, f"Resource MIME type should be valid: {resource.mime_type}" + + def test_tool_configuration_consistency(self): + """Test tool configuration consistency.""" + from codegen.cli.mcp.server import mcp + + # All tools should have names and functions + tools = mcp._tool_manager._tools + for name, tool in tools.items(): + assert hasattr(tool, "fn"), "Tool should have function" + + # Name should be a string + assert isinstance(name, str), "Tool name should be string" + assert len(name) > 0, "Tool name should not be empty" + + # Function should be callable + assert callable(tool.fn), "Tool function should be callable" diff --git a/tests/cli/mcp/test_mcp_error_handling.py b/tests/cli/mcp/test_mcp_error_handling.py new file mode 100644 index 000000000..3bb6458c2 --- /dev/null +++ b/tests/cli/mcp/test_mcp_error_handling.py @@ -0,0 +1,374 @@ +"""Tests for MCP error handling scenarios.""" + +import os +import subprocess +import time +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest + + +class TestMCPErrorHandling: + """Test MCP error handling scenarios.""" + + def test_mcp_server_startup_without_dependencies(self): + """Test MCP server behavior when optional dependencies are missing.""" + codegen_path = Path(__file__).parent.parent.parent.parent / "src" + venv_python = Path(__file__).parent.parent.parent.parent / ".venv" / "bin" / "python" + + env = os.environ.copy() + env["PYTHONPATH"] = str(codegen_path) + # Remove API key to test behavior without authentication + env.pop("CODEGEN_API_KEY", None) + + process = subprocess.Popen( + [str(venv_python), "-c", "from codegen.cli.cli import main; main(['mcp', '--transport', 'stdio'])"], + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + try: + # Give it time to start + time.sleep(2) + + # Server should start even without API key + assert process.poll() is None, "Server should start without API key" + + finally: + # Cleanup + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + def test_mcp_server_with_invalid_api_base_url(self): + """Test MCP server behavior with invalid API base URL.""" + codegen_path = Path(__file__).parent.parent.parent.parent / "src" + venv_python = Path(__file__).parent.parent.parent.parent / ".venv" / "bin" / "python" + + env = os.environ.copy() + env["PYTHONPATH"] = str(codegen_path) + env["CODEGEN_API_KEY"] = "test-key" + env["CODEGEN_API_BASE_URL"] = "invalid-url" + + process = subprocess.Popen( + [str(venv_python), "-c", "from codegen.cli.cli import main; main(['mcp', '--transport', 'stdio'])"], + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + try: + # Give it time to start + time.sleep(2) + + # Server should start even with invalid API URL + assert process.poll() is None, "Server should start with invalid API URL" + + finally: + # Cleanup + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + @patch("codegen.cli.mcp.server.API_CLIENT_AVAILABLE", False) + def test_api_client_unavailable_error_handling(self): + """Test error handling when API client is not available.""" + from codegen.cli.mcp.server import get_api_client + + with pytest.raises(RuntimeError, match="codegen-api-client is not available"): + get_api_client() + + @patch("codegen.cli.mcp.server.get_api_client") + def test_create_agent_run_api_error_handling(self, mock_get_api_client): + """Test create_agent_run tool error handling.""" + from codegen.cli.mcp.server import mcp + + # Mock API to raise an exception + mock_get_api_client.side_effect = Exception("Network error") + + # Get the create_agent_run tool function + tools = mcp._tool_manager._tools + assert "create_agent_run" in tools + + create_agent_run_tool = tools["create_agent_run"] + + # Test the tool function with API error + result = create_agent_run_tool.fn( # type: ignore[attr-defined] + org_id=1, + prompt="Test prompt", + ctx=None + ) + + + assert "Error creating agent run" in result + assert "Network error" in result + + @patch("codegen.cli.mcp.server.get_api_client") + def test_get_agent_run_api_error_handling(self, mock_get_api_client): + """Test get_agent_run tool error handling.""" + from codegen.cli.mcp.server import mcp + + # Mock API to raise an exception + mock_get_api_client.side_effect = Exception("API timeout") + + # Get the get_agent_run tool function + tools = mcp._tool_manager._tools + assert "get_agent_run" in tools + get_agent_run_tool = tools["get_agent_run"] + + get_agent_run_tool = tools["get_agent_run"] + + # Test the tool function with API error + result = get_agent_run_tool.fn( # type: ignore[attr-defined] + org_id=1, + agent_run_id=123, + ctx=None + ) + + + assert "Error getting agent run" in result + assert "API timeout" in result + + @patch("codegen.cli.mcp.server.get_api_client") + def test_get_organizations_api_error_handling(self, mock_get_api_client): + """Test get_organizations tool error handling.""" + from codegen.cli.mcp.server import mcp + + # Mock API to raise an exception + mock_get_api_client.side_effect = Exception("Authentication failed") + + # Get the get_organizations tool function + tools = mcp._tool_manager._tools + assert "get_organizations" in tools + get_organizations_tool = tools["get_organizations"] + + get_organizations_tool = tools["get_organizations"] + + # Test the tool function with API error + result = get_organizations_tool.fn( # type: ignore[attr-defined] + page=1, limit=10, ctx=None) + + + assert "Error getting organizations" in result + assert "Authentication failed" in result + + @patch("codegen.cli.mcp.server.get_api_client") + def test_get_users_api_error_handling(self, mock_get_api_client): + """Test get_users tool error handling.""" + from codegen.cli.mcp.server import mcp + + # Mock API to raise an exception + mock_get_api_client.side_effect = Exception("Permission denied") + + # Get the get_users tool function + tools = mcp._tool_manager._tools + assert "get_users" in tools + get_users_tool = tools["get_users"] + + get_users_tool = tools["get_users"] + + # Test the tool function with API error + result = get_users_tool.fn( # type: ignore[attr-defined] + org_id=1, page=1, limit=10, ctx=None) + + + assert "Error getting users" in result + assert "Permission denied" in result + + @patch("codegen.cli.mcp.server.get_api_client") + def test_get_user_api_error_handling(self, mock_get_api_client): + """Test get_user tool error handling.""" + from codegen.cli.mcp.server import mcp + + # Mock API to raise an exception + mock_get_api_client.side_effect = Exception("User not found") + + # Get the get_user tool function + tools = mcp._tool_manager._tools + assert "get_user" in tools + get_user_tool = tools["get_user"] + + get_user_tool = tools["get_user"] + + # Test the tool function with API error + result = get_user_tool.fn( # type: ignore[attr-defined] + org_id=1, user_id=999, ctx=None) + + + assert "Error getting user" in result + assert "User not found" in result + + def test_run_server_invalid_transport_error(self): + """Test run_server function with invalid transport.""" + from codegen.cli.mcp.server import run_server + + with pytest.raises(ValueError, match="Unsupported transport: invalid"): + run_server(transport="invalid") + + def test_run_server_http_transport_fallback(self): + """Test run_server function HTTP transport fallback.""" + from codegen.cli.mcp.server import run_server + + # This should not raise an exception but fall back to stdio + # We can't easily test the actual behavior without mocking FastMCP + # So we'll just ensure it doesn't crash + try: + # This will actually try to run the server, so we need to be careful + # For now, we'll just test that the function exists and can be called + assert callable(run_server) + except Exception: + # If it raises an exception, it should be a controlled one + pass + + @patch("codegen.cli.mcp.server.LEGACY_IMPORTS_AVAILABLE", False) + def test_legacy_tools_unavailable(self): + """Test behavior when legacy imports are not available.""" + # Re-import the server module to trigger the conditional logic + import importlib + + import codegen.cli.mcp.server + + importlib.reload(codegen.cli.mcp.server) + + from codegen.cli.mcp.server import mcp + + # Check that legacy tools are not registered when imports are unavailable + tool_names = [tool.name for tool in mcp._tools] + + # ask_codegen_sdk and improve_codemod should not be available + legacy_tools = ["ask_codegen_sdk", "improve_codemod"] + for legacy_tool in legacy_tools: + assert legacy_tool not in tool_names, f"Legacy tool {legacy_tool} should not be available" + + def test_api_client_configuration_error_handling(self): + """Test API client configuration error handling.""" + from codegen.cli.mcp.server import get_api_client + + # Test with missing environment variables + with patch.dict(os.environ, {}, clear=True): + try: + get_api_client() + # Should not raise an exception, but should use defaults + except Exception as e: + # If it does raise an exception, it should be a controlled one + assert "codegen-api-client is not available" in str(e) + + @patch("codegen.cli.mcp.server.get_api_client") + def test_api_client_initialization_error(self, mock_get_api_client): + """Test API client initialization error handling.""" + from codegen.cli.mcp.server import mcp + + # Mock API client initialization to fail + mock_get_api_client.side_effect = RuntimeError("Failed to initialize API client") + + # Get any API-dependent tool + tools = mcp._tool_manager._tools + assert "create_agent_run" in tools + create_agent_run_tool = tools["create_agent_run"] + + + tool = list(tools.values())[0] + + # Test that the tool handles initialization errors gracefully + result = tool.fn( # type: ignore[attr-defined] + org_id=1, prompt="test", ctx=None) + + + assert "Error creating agent run" in result + assert "Failed to initialize API client" in result + + def test_resource_import_error_handling(self): + """Test resource import error handling.""" + # Test that resources can be imported without errors + try: + from codegen.cli.mcp.resources.system_prompt import SYSTEM_PROMPT + from codegen.cli.mcp.resources.system_setup_instructions import SETUP_INSTRUCTIONS + + assert isinstance(SYSTEM_PROMPT, str) + assert isinstance(SETUP_INSTRUCTIONS, str) + except ImportError as e: + pytest.fail( # type: ignore[misc] + f"Resource imports should not fail: {e}") + + def test_server_module_import_error_handling(self): + """Test server module import error handling.""" + # Test that the server module can be imported without errors + try: + import codegen.cli.mcp.server + + assert hasattr(codegen.cli.mcp.server, "run_server") + assert hasattr(codegen.cli.mcp.server, "mcp") + except ImportError as e: + pytest.fail( # type: ignore[misc] + f"Server module import should not fail: {e}") + + def test_mcp_command_import_error_handling(self): + """Test MCP command import error handling.""" + # Test that the MCP command can be imported without errors + try: + from codegen.cli.commands.mcp.main import mcp + + assert callable(mcp) + except ImportError as e: + pytest.fail( # type: ignore[misc] + f"MCP command import should not fail: {e}") + + @patch("codegen.cli.mcp.server.get_api_client") + def test_api_response_parsing_error(self, mock_get_api_client): + """Test API response parsing error handling.""" + from codegen.cli.mcp.server import mcp + + # Mock API to return malformed response + mock_response = Mock() + mock_response.id = None # This could cause issues + mock_response.status = None + mock_response.created_at = None + mock_response.prompt = None + mock_response.repo_name = None + mock_response.branch_name = None + + mock_agents_api = Mock() + mock_agents_api.create_agent_run_v1_organizations_org_id_agent_run_post.return_value = mock_response + + mock_get_api_client.return_value = (None, mock_agents_api, None, None) + + # Get the create_agent_run tool function + tools = mcp._tool_manager._tools + assert "create_agent_run" in tools + create_agent_run_tool = tools["create_agent_run"] + + create_agent_run_tool = tools["create_agent_run"] + + # Test the tool function with malformed response + result = create_agent_run_tool.fn( # type: ignore[attr-defined] + org_id=1, + prompt="Test prompt", + ctx=None + ) + + + # Should handle None values gracefully + assert isinstance(result, str) + # Should be valid JSON even with None values + import json + + try: + parsed = json.loads(result) + assert isinstance(parsed, dict) + except json.JSONDecodeError: + pytest.fail( # type: ignore[misc] + "Tool should return valid JSON even with malformed API response") diff --git a/tests/cli/mcp/test_mcp_protocol.py b/tests/cli/mcp/test_mcp_protocol.py new file mode 100644 index 000000000..5637d6992 --- /dev/null +++ b/tests/cli/mcp/test_mcp_protocol.py @@ -0,0 +1,254 @@ +"""Tests for MCP protocol communication.""" + +import json +import os +import subprocess +import time +from pathlib import Path +from unittest.mock import patch + +import pytest + + +class TestMCPProtocol: + """Test MCP protocol communication.""" + + @pytest.fixture + def mcp_server_process(self): + """Start an MCP server process for testing.""" + codegen_path = Path(__file__).parent.parent.parent.parent / "src" + venv_python = Path(__file__).parent.parent.parent.parent / ".venv" / "bin" / "python" + + env = os.environ.copy() + env["PYTHONPATH"] = str(codegen_path) + # Set mock API key for testing + env["CODEGEN_API_KEY"] = "test-api-key" + env["CODEGEN_API_BASE_URL"] = "https://api.test.codegen.com" + + process = subprocess.Popen( + [str(venv_python), "-c", "from codegen.cli.cli import main; main(['mcp', '--transport', 'stdio'])"], + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Give it time to start + time.sleep(2) + + yield process + + # Cleanup + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + def send_mcp_message(self, process, message): + """Send an MCP message to the server.""" + if process.stdin: + process.stdin.write(json.dumps(message) + "\n") + process.stdin.flush() + time.sleep(0.5) + + def test_mcp_initialization_sequence(self, mcp_server_process): + """Test the complete MCP initialization sequence.""" + # Send initialize message + init_message = { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {"protocolVersion": "2024-11-05", "capabilities": {"roots": {"listChanged": True}, "sampling": {}}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}, + } + + self.send_mcp_message(mcp_server_process, init_message) + + # Server should still be running after initialization + assert mcp_server_process.poll() is None + + # Send initialized notification + initialized_message = {"jsonrpc": "2.0", "method": "notifications/initialized"} + + self.send_mcp_message(mcp_server_process, initialized_message) + + # Server should still be running + assert mcp_server_process.poll() is None + + def test_mcp_list_resources_request(self, mcp_server_process): + """Test MCP resources/list request.""" + # Initialize first + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + self.send_mcp_message(mcp_server_process, init_message) + + # List resources + list_resources_message = {"jsonrpc": "2.0", "id": 2, "method": "resources/list"} + + self.send_mcp_message(mcp_server_process, list_resources_message) + + # Server should handle the request without crashing + assert mcp_server_process.poll() is None + + def test_mcp_list_tools_request(self, mcp_server_process): + """Test MCP tools/list request.""" + # Initialize first + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + self.send_mcp_message(mcp_server_process, init_message) + + # List tools + list_tools_message = {"jsonrpc": "2.0", "id": 3, "method": "tools/list"} + + self.send_mcp_message(mcp_server_process, list_tools_message) + + # Server should handle the request without crashing + assert mcp_server_process.poll() is None + + def test_mcp_resource_read_request(self, mcp_server_process): + """Test MCP resources/read request.""" + # Initialize first + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + self.send_mcp_message(mcp_server_process, init_message) + + # Read a resource + read_resource_message = {"jsonrpc": "2.0", "id": 4, "method": "resources/read", "params": {"uri": "system://agent_prompt"}} + + self.send_mcp_message(mcp_server_process, read_resource_message) + + # Server should handle the request without crashing + assert mcp_server_process.poll() is None + + @patch("codegen.cli.mcp.server.get_api_client") + def test_mcp_tool_call_request(self, mock_get_api_client, mcp_server_process): + """Test MCP tools/call request.""" + # Mock API client to avoid actual API calls + mock_get_api_client.return_value = (None, None, None, None) + + # Initialize first + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + self.send_mcp_message(mcp_server_process, init_message) + + # Call a tool + tool_call_message = { + "jsonrpc": "2.0", + "id": 5, + "method": "tools/call", + "params": {"name": "generate_codemod", "arguments": {"title": "test-codemod", "task": "Add logging to functions", "codebase_path": "/tmp/test"}}, + } + + self.send_mcp_message(mcp_server_process, tool_call_message) + + # Server should handle the request without crashing + assert mcp_server_process.poll() is None + + def test_mcp_invalid_method_request(self, mcp_server_process): + """Test MCP server handling of invalid method requests.""" + # Initialize first + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + self.send_mcp_message(mcp_server_process, init_message) + + # Send invalid method + invalid_message = {"jsonrpc": "2.0", "id": 6, "method": "invalid/method"} + + self.send_mcp_message(mcp_server_process, invalid_message) + + # Server should handle the invalid request gracefully + assert mcp_server_process.poll() is None + + def test_mcp_malformed_json_handling(self, mcp_server_process): + """Test MCP server handling of malformed JSON.""" + # Send malformed JSON + if mcp_server_process.stdin: + mcp_server_process.stdin.write("{ invalid json }\n") + mcp_server_process.stdin.flush() + + time.sleep(0.5) + + # Server should handle malformed JSON gracefully + assert mcp_server_process.poll() is None + + # Server should still be able to handle valid requests after malformed JSON + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + self.send_mcp_message(mcp_server_process, init_message) + + # Server should still be running + assert mcp_server_process.poll() is None + + def test_mcp_missing_required_fields(self, mcp_server_process): + """Test MCP server handling of messages with missing required fields.""" + # Send message without required fields + incomplete_message = { + "jsonrpc": "2.0", + "method": "initialize", + # Missing id and params + } + + self.send_mcp_message(mcp_server_process, incomplete_message) + + # Server should handle incomplete messages gracefully + assert mcp_server_process.poll() is None + + def test_mcp_protocol_version_handling(self, mcp_server_process): + """Test MCP server handling of different protocol versions.""" + # Test with older protocol version + init_message_old = { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2024-01-01", # Older version + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.0.0"}, + }, + } + + self.send_mcp_message(mcp_server_process, init_message_old) + + # Server should handle different protocol versions + assert mcp_server_process.poll() is None + + def test_mcp_concurrent_requests(self, mcp_server_process): + """Test MCP server handling of concurrent requests.""" + # Initialize first + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + self.send_mcp_message(mcp_server_process, init_message) + + # Send multiple requests quickly + messages = [ + {"jsonrpc": "2.0", "id": 2, "method": "resources/list"}, + {"jsonrpc": "2.0", "id": 3, "method": "tools/list"}, + {"jsonrpc": "2.0", "id": 4, "method": "resources/read", "params": {"uri": "system://manifest"}}, + ] + + for message in messages: + self.send_mcp_message(mcp_server_process, message) + + # Give time for all requests to be processed + time.sleep(2) + + # Server should handle concurrent requests without crashing + assert mcp_server_process.poll() is None + + def test_mcp_server_shutdown_handling(self, mcp_server_process): + """Test MCP server graceful shutdown.""" + # Initialize first + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + self.send_mcp_message(mcp_server_process, init_message) + + # Send shutdown request + shutdown_message = {"jsonrpc": "2.0", "id": 99, "method": "shutdown"} + + self.send_mcp_message(mcp_server_process, shutdown_message) + + # Send exit notification + exit_message = {"jsonrpc": "2.0", "method": "exit"} + + self.send_mcp_message(mcp_server_process, exit_message) + + # Give time for graceful shutdown + time.sleep(2) + + # Server should have shut down gracefully + # Note: This test might be flaky depending on FastMCP's shutdown handling diff --git a/tests/cli/mcp/test_mcp_resources.py b/tests/cli/mcp/test_mcp_resources.py new file mode 100644 index 000000000..a68efd9c8 --- /dev/null +++ b/tests/cli/mcp/test_mcp_resources.py @@ -0,0 +1,221 @@ +"""Tests for MCP resources functionality.""" + +import json +from unittest.mock import patch + +import pytest + + +class TestMCPResources: + """Test MCP resources functionality.""" + + def test_system_agent_prompt_resource(self): + """Test the system://agent_prompt resource.""" + from codegen.cli.mcp.server import mcp + + # Get the agent_prompt resource function + resources = mcp._resource_manager._resources + assert "system://agent_prompt" in resources + + agent_prompt_resource = resources["system://agent_prompt"] + + # Test the resource function + result = agent_prompt_resource.fn( # type: ignore[attr-defined] + ) + + + # Should return a string with system prompt content + assert isinstance(result, str) + assert len(result) > 0 + # Should contain information about Codegen SDK + assert "codegen" in result.lower() or "sdk" in result.lower() + + def test_system_setup_instructions_resource(self): + """Test the system://setup_instructions resource.""" + from codegen.cli.mcp.server import mcp + + # Get the setup_instructions resource function + resources = mcp._resource_manager._resources + assert "system://setup_instructions" in resources + + setup_instructions_resource = resources["system://setup_instructions"] + + # Test the resource function + result = setup_instructions_resource.fn( # type: ignore[attr-defined] + ) + + + # Should return a string with setup instructions + assert isinstance(result, str) + assert len(result) > 0 + # Should contain setup-related content + assert any(keyword in result.lower() for keyword in ["setup", "install", "environment", "configure"]) + + def test_system_manifest_resource(self): + """Test the system://manifest resource.""" + from codegen.cli.mcp.server import mcp + + # Get the manifest resource function + resources = mcp._resource_manager._resources + assert "system://manifest" in resources + + manifest_resource = resources["system://manifest"] + + # Test the resource function + result = manifest_resource.fn( # type: ignore[attr-defined] + ) + + + # Should return a dictionary with manifest information + assert isinstance(result, dict) + assert "name" in result + assert "version" in result + assert "description" in result + + # Check specific values + assert result["name"] == "mcp-codegen" + assert result["version"] == "0.1.0" + assert "codegen" in result["description"].lower() + + def test_resource_mime_types(self): + """Test that resources have correct MIME types.""" + from codegen.cli.mcp.server import mcp + + resources = mcp._resource_manager._resources + resource_mime_types = {uri: resource.mime_type for uri, resource in resources.items()} + + # Check expected MIME types + assert resource_mime_types["system://agent_prompt"] == "text/plain" + assert resource_mime_types["system://setup_instructions"] == "text/plain" + assert resource_mime_types["system://manifest"] == "application/json" + + def test_resource_descriptions(self): + """Test that resources have appropriate descriptions.""" + from codegen.cli.mcp.server import mcp + + resources = mcp._resource_manager._resources + resource_descriptions = {uri: resource.description for uri, resource in resources.items()} + + # Check that descriptions exist and are meaningful + assert "agent" in resource_descriptions["system://agent_prompt"].lower() + assert "codegen" in resource_descriptions["system://agent_prompt"].lower() + + assert "setup" in resource_descriptions["system://setup_instructions"].lower() + assert "instructions" in resource_descriptions["system://setup_instructions"].lower() + + # Manifest resource might not have a description, but if it does, it should be meaningful + if resource_descriptions.get("system://manifest"): + assert len(resource_descriptions["system://manifest"]) > 0 + + def test_all_resources_callable(self): + """Test that all resource functions are callable.""" + from codegen.cli.mcp.server import mcp + + resources = mcp._resource_manager._resources + for uri, resource in resources.items(): + assert callable(resource.fn), f"Resource {uri} function is not callable" + + # Try calling the function + try: + result = resource.fn( # type: ignore[attr-defined] + ) + assert result is not None, f"Resource {uri} returned None" + except Exception as e: + pytest.fail( # type: ignore[misc] + f"Resource {uri} raised exception: {e}") + + def test_resource_content_consistency(self): + """Test that resource content is consistent across calls.""" + from codegen.cli.mcp.server import mcp + + resources = mcp._resource_manager._resources + for uri, resource in resources.items(): + # Call the resource function multiple times + result1 = resource.fn( # type: ignore[attr-defined] + ) + result2 = resource.fn( # type: ignore[attr-defined] + ) + + + # Results should be identical (resources should be deterministic) + assert result1 == result2, f"Resource {uri} returned different results on multiple calls" + + def test_system_prompt_content_structure(self): + """Test that the system prompt has expected structure.""" + from codegen.cli.mcp.resources.system_prompt import SYSTEM_PROMPT + + # Should be a non-empty string + assert isinstance(SYSTEM_PROMPT, str) + assert len(SYSTEM_PROMPT) > 100 # Should be substantial + + # Should contain key information about Codegen + assert "codegen" in SYSTEM_PROMPT.lower() + + def test_setup_instructions_content_structure(self): + """Test that setup instructions have expected structure.""" + from codegen.cli.mcp.resources.system_setup_instructions import SETUP_INSTRUCTIONS + + # Should be a non-empty string + assert isinstance(SETUP_INSTRUCTIONS, str) + assert len(SETUP_INSTRUCTIONS) > 50 # Should be substantial + + # Should contain setup-related keywords + setup_keywords = ["install", "setup", "configure", "environment", "python", "pip", "uv"] + assert any(keyword in SETUP_INSTRUCTIONS.lower() for keyword in setup_keywords) + + @patch("codegen.cli.mcp.server.mcp") + def test_resource_registration_process(self, mock_mcp): + """Test that resources are properly registered with the MCP server.""" + # Import the server module to trigger resource registration + + # Check that the resource decorator was called + assert mock_mcp.resource.called, "MCP resource decorator should have been called" + + # Check that it was called the expected number of times + expected_resource_count = 3 # agent_prompt, setup_instructions, manifest + assert mock_mcp.resource.call_count >= expected_resource_count + + def test_resource_error_handling(self): + """Test resource error handling.""" + from codegen.cli.mcp.server import mcp + + # All current resources should not raise exceptions + resources = mcp._resource_manager._resources + for uri, resource in resources.items(): + try: + result = resource.fn( # type: ignore[attr-defined] + ) + # Basic validation that result is not empty + if isinstance(result, str): + assert len(result) > 0 + elif isinstance(result, dict): + assert len(result) > 0 + else: + pytest.fail( # type: ignore[misc] + f"Resource {uri} returned unexpected type: {type(result)}") + except Exception as e: + pytest.fail( # type: ignore[misc] + f"Resource {uri} should not raise exceptions, but raised: {e}") + + def test_json_serializable_manifest(self): + """Test that the manifest resource returns JSON-serializable data.""" + from codegen.cli.mcp.server import mcp + + # Get the manifest resource + manifest_resources = mcp._resource_manager._resources + assert "system://manifest" in manifest_resources + + manifest_resource = manifest_resources["system://manifest"] + result = manifest_resource.fn( # type: ignore[attr-defined] + ) + + + # Should be JSON serializable + try: + json_str = json.dumps(result) + # Should be able to parse it back + parsed = json.loads(json_str) + assert parsed == result + except (TypeError, ValueError) as e: + pytest.fail( # type: ignore[misc] + f"Manifest resource result is not JSON serializable: {e}") diff --git a/tests/cli/mcp/test_mcp_tools.py b/tests/cli/mcp/test_mcp_tools.py new file mode 100644 index 000000000..8d077e533 --- /dev/null +++ b/tests/cli/mcp/test_mcp_tools.py @@ -0,0 +1,397 @@ +"""Comprehensive tests for MCP tools with mocked API calls.""" + +import json +import os +import subprocess +import time +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest + + +class TestMCPTools: + """Test MCP tools functionality with mocked API calls.""" + + @pytest.fixture + def mcp_server_process(self): + """Start an MCP server process for testing.""" + codegen_path = Path(__file__).parent.parent.parent.parent / "src" + venv_python = Path(__file__).parent.parent.parent.parent / ".venv" / "bin" / "python" + + env = os.environ.copy() + env["PYTHONPATH"] = str(codegen_path) + # Set mock API key for testing + env["CODEGEN_API_KEY"] = "test-api-key" + env["CODEGEN_API_BASE_URL"] = "https://api.test.codegen.com" + + process = subprocess.Popen( + [str(venv_python), "-c", "from codegen.cli.cli import main; main(['mcp', '--transport', 'stdio'])"], + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + # Give it time to start + time.sleep(2) + + yield process + + # Cleanup + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + def send_mcp_message(self, process, message): + """Send an MCP message to the server and get response.""" + if process.stdin: + process.stdin.write(json.dumps(message) + "\n") + process.stdin.flush() + + # Give time to process + time.sleep(0.5) + + # Read response (this is simplified - real MCP would need proper parsing) + return None # For now, we'll test that the server doesn't crash + + def test_mcp_initialize(self, mcp_server_process): + """Test MCP server initialization.""" + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + + self.send_mcp_message(mcp_server_process, init_message) + + # Server should still be running after initialization + assert mcp_server_process.poll() is None + + def test_mcp_list_resources(self, mcp_server_process): + """Test listing MCP resources.""" + # First initialize + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + self.send_mcp_message(mcp_server_process, init_message) + + # List resources + list_resources_message = {"jsonrpc": "2.0", "id": 2, "method": "resources/list"} + + self.send_mcp_message(mcp_server_process, list_resources_message) + + # Server should still be running + assert mcp_server_process.poll() is None + + def test_mcp_list_tools(self, mcp_server_process): + """Test listing MCP tools.""" + # First initialize + init_message = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "test-client", "version": "1.0.0"}}} + self.send_mcp_message(mcp_server_process, init_message) + + # List tools + list_tools_message = {"jsonrpc": "2.0", "id": 3, "method": "tools/list"} + + self.send_mcp_message(mcp_server_process, list_tools_message) + + # Server should still be running + assert mcp_server_process.poll() is None + + @patch("codegen.cli.mcp.server.get_api_client") + def test_generate_codemod_tool(self, mock_get_api_client): + """Test the generate_codemod tool.""" + from codegen.cli.mcp.server import mcp + + # Get the generate_codemod tool function + tools = mcp._tool_manager._tools + assert "generate_codemod" in tools + + generate_codemod_tool = tools["generate_codemod"] + + # Test the tool function + result = generate_codemod_tool.fn( # type: ignore[attr-defined] + title="test-codemod", + task="Add logging to all functions", + codebase_path="/path/to/codebase", + ctx=None + ) + + + assert "codegen create test-codemod" in result + assert "Add logging to all functions" in result + + @patch("codegen.cli.mcp.server.get_api_client") + def test_create_agent_run_tool_success(self, mock_get_api_client): + """Test the create_agent_run tool with successful API response.""" + from codegen.cli.mcp.server import mcp + + # Mock the API response + mock_response = Mock() + mock_response.id = 123 + mock_response.status = "running" + mock_response.created_at = datetime.now(timezone.utc) + mock_response.prompt = "Test prompt" + mock_response.repo_name = "test-repo" + mock_response.branch_name = "test-branch" + + mock_agents_api = Mock() + mock_agents_api.create_agent_run_v1_organizations_org_id_agent_run_post.return_value = mock_response + + mock_get_api_client.return_value = (None, mock_agents_api, None, None) + + # Get the create_agent_run tool function + tools = mcp._tool_manager._tools + assert "create_agent_run" in tools + + create_agent_run_tool = tools["create_agent_run"] + + # Test the tool function + result = create_agent_run_tool.fn( # type: ignore[attr-defined] + org_id=1, + prompt="Test prompt", + repo_name="test-repo", + branch_name="test-branch", + ctx=None + ) + + + # Parse the JSON response + response_data = json.loads(result) + assert response_data["id"] == 123 + assert response_data["status"] == "running" + assert response_data["prompt"] == "Test prompt" + assert response_data["repo_name"] == "test-repo" + assert response_data["branch_name"] == "test-branch" + + @patch("codegen.cli.mcp.server.get_api_client") + def test_create_agent_run_tool_error(self, mock_get_api_client): + """Test the create_agent_run tool with API error.""" + from codegen.cli.mcp.server import mcp + + # Mock API to raise an exception + mock_get_api_client.side_effect = Exception("API connection failed") + + # Get the create_agent_run tool function + tools = mcp._tool_manager._tools + assert "create_agent_run" in tools + create_agent_run_tool = tools["create_agent_run"] + + create_agent_run_tool = tools["create_agent_run"] + + # Test the tool function + result = create_agent_run_tool.fn( # type: ignore[attr-defined] + org_id=1, + prompt="Test prompt", + ctx=None + ) + + + assert "Error creating agent run" in result + assert "API connection failed" in result + + @patch("codegen.cli.mcp.server.get_api_client") + def test_get_agent_run_tool_success(self, mock_get_api_client): + """Test the get_agent_run tool with successful API response.""" + from codegen.cli.mcp.server import mcp + + # Mock the API response + mock_response = Mock() + mock_response.id = 123 + mock_response.status = "completed" + mock_response.created_at = datetime.now(timezone.utc) + mock_response.updated_at = datetime.now(timezone.utc) + mock_response.prompt = "Test prompt" + mock_response.repo_name = "test-repo" + mock_response.branch_name = "test-branch" + mock_response.result = "Task completed successfully" + + mock_agents_api = Mock() + mock_agents_api.get_agent_run_v1_organizations_org_id_agent_run_agent_run_id_get.return_value = mock_response + + mock_get_api_client.return_value = (None, mock_agents_api, None, None) + + # Get the get_agent_run tool function + tools = mcp._tool_manager._tools + assert "get_agent_run" in tools + get_agent_run_tool = tools["get_agent_run"] + + get_agent_run_tool = tools["get_agent_run"] + + # Test the tool function + result = get_agent_run_tool.fn( # type: ignore[attr-defined] + org_id=1, + agent_run_id=123, + ctx=None + ) + + + # Parse the JSON response + response_data = json.loads(result) + assert response_data["id"] == 123 + assert response_data["status"] == "completed" + assert response_data["result"] == "Task completed successfully" + + @patch("codegen.cli.mcp.server.get_api_client") + def test_get_organizations_tool_success(self, mock_get_api_client): + """Test the get_organizations tool with successful API response.""" + from datetime import datetime + + from codegen.cli.mcp.server import mcp + + # Mock organization objects + mock_org1 = Mock() + mock_org1.id = 1 + mock_org1.name = "Test Org 1" + mock_org1.slug = "test-org-1" + mock_org1.created_at = datetime.now(timezone.utc) + + mock_org2 = Mock() + mock_org2.id = 2 + mock_org2.name = "Test Org 2" + mock_org2.slug = "test-org-2" + mock_org2.created_at = datetime.now(timezone.utc) + + # Mock the API response + mock_response = Mock() + mock_response.items = [mock_org1, mock_org2] + mock_response.total = 2 + mock_response.page = 1 + mock_response.limit = 10 + + mock_organizations_api = Mock() + mock_organizations_api.get_organizations_v1_organizations_get.return_value = mock_response + + mock_get_api_client.return_value = (None, None, mock_organizations_api, None) + + # Get the get_organizations tool function + tools = mcp._tool_manager._tools + assert "get_organizations" in tools + get_organizations_tool = tools["get_organizations"] + + get_organizations_tool = tools["get_organizations"] + + # Test the tool function + result = get_organizations_tool.fn( # type: ignore[attr-defined] + page=1, limit=10, ctx=None) + + + # Parse the JSON response + response_data = json.loads(result) + assert response_data["total"] == 2 + assert len(response_data["organizations"]) == 2 + assert response_data["organizations"][0]["name"] == "Test Org 1" + assert response_data["organizations"][1]["name"] == "Test Org 2" + + @patch("codegen.cli.mcp.server.get_api_client") + def test_get_users_tool_success(self, mock_get_api_client): + """Test the get_users tool with successful API response.""" + from datetime import datetime + + from codegen.cli.mcp.server import mcp + + # Mock user objects + mock_user1 = Mock() + mock_user1.id = 1 + mock_user1.email = "user1@test.com" + mock_user1.name = "User One" + mock_user1.created_at = datetime.now(timezone.utc) + + mock_user2 = Mock() + mock_user2.id = 2 + mock_user2.email = "user2@test.com" + mock_user2.name = "User Two" + mock_user2.created_at = datetime.now(timezone.utc) + + # Mock the API response + mock_response = Mock() + mock_response.items = [mock_user1, mock_user2] + mock_response.total = 2 + mock_response.page = 1 + mock_response.limit = 10 + + mock_users_api = Mock() + mock_users_api.get_users_v1_organizations_org_id_users_get.return_value = mock_response + + mock_get_api_client.return_value = (None, None, None, mock_users_api) + + # Get the get_users tool function + tools = mcp._tool_manager._tools + assert "get_users" in tools + get_users_tool = tools["get_users"] + + get_users_tool = tools["get_users"] + + # Test the tool function + result = get_users_tool.fn( # type: ignore[attr-defined] + org_id=1, page=1, limit=10, ctx=None) + + + # Parse the JSON response + response_data = json.loads(result) + assert response_data["total"] == 2 + assert len(response_data["users"]) == 2 + assert response_data["users"][0]["email"] == "user1@test.com" + assert response_data["users"][1]["email"] == "user2@test.com" + + @patch("codegen.cli.mcp.server.get_api_client") + def test_get_user_tool_success(self, mock_get_api_client): + """Test the get_user tool with successful API response.""" + from datetime import datetime + + from codegen.cli.mcp.server import mcp + + # Mock the API response + mock_response = Mock() + mock_response.id = 1 + mock_response.email = "user@test.com" + mock_response.name = "Test User" + mock_response.created_at = datetime.now(timezone.utc) + mock_response.updated_at = datetime.now(timezone.utc) + + mock_users_api = Mock() + mock_users_api.get_user_v1_organizations_org_id_users_user_id_get.return_value = mock_response + + mock_get_api_client.return_value = (None, None, None, mock_users_api) + + # Get the get_user tool function + tools = mcp._tool_manager._tools + assert "get_user" in tools + get_user_tool = tools["get_user"] + + get_user_tool = tools["get_user"] + + # Test the tool function + result = get_user_tool.fn( # type: ignore[attr-defined] + org_id=1, user_id=1, ctx=None) + + + # Parse the JSON response + response_data = json.loads(result) + assert response_data["id"] == 1 + assert response_data["email"] == "user@test.com" + assert response_data["name"] == "Test User" + + def test_mcp_tools_registration(self): + """Test that all expected tools are registered.""" + from codegen.cli.mcp.server import mcp + + tool_names = list(mcp._tool_manager._tools.keys()) + + # Check that all expected tools are registered + expected_tools = ["generate_codemod", "create_agent_run", "get_agent_run", "get_organizations", "get_users", "get_user"] + + for expected_tool in expected_tools: + assert expected_tool in tool_names, f"Tool {expected_tool} not found in registered tools" + + def test_mcp_resources_registration(self): + """Test that all expected resources are registered.""" + from codegen.cli.mcp.server import mcp + + resource_uris = list(mcp._resource_manager._resources.keys()) + + # Check that all expected resources are registered + expected_resources = ["system://agent_prompt", "system://setup_instructions", "system://manifest"] + + for expected_resource in expected_resources: + assert expected_resource in resource_uris, f"Resource {expected_resource} not found in registered resources" diff --git a/tests/cli/mcp/test_simple_integration.py b/tests/cli/mcp/test_simple_integration.py index 86dac31e1..54bc5b2fb 100644 --- a/tests/cli/mcp/test_simple_integration.py +++ b/tests/cli/mcp/test_simple_integration.py @@ -2,8 +2,6 @@ from unittest.mock import patch -import pytest - class TestMCPSimpleIntegration: """Simple integration tests that avoid FastMCP import issues.""" @@ -22,7 +20,8 @@ def test_api_client_imports_available(self): assert UsersApi is not None except ImportError as e: - raise AssertionError(f"API client imports not available: {e}") from e + msg = f"API client imports not available: {e}" + raise AssertionError(msg) from e def test_mcp_command_registration(self): """Test that the MCP command is registered in the CLI.""" @@ -58,7 +57,8 @@ def test_server_configuration_basic(self): assert callable(mcp) except ImportError as e: - raise AssertionError(f"MCP command module not importable: {e}") from e + msg = f"MCP command module not importable: {e}" + raise AssertionError(msg) from e def test_environment_variable_handling_basic(self): """Test basic environment variable handling.""" diff --git a/tests/unit/codegen/agents/test_agent.py b/tests/unit/codegen/agents/test_agent.py index bc4e448ba..2c994299f 100644 --- a/tests/unit/codegen/agents/test_agent.py +++ b/tests/unit/codegen/agents/test_agent.py @@ -73,7 +73,7 @@ def test_agent_init_requires_token(self): """Test that Agent initialization requires a token parameter.""" # This should raise a TypeError because token is a required parameter with pytest.raises(TypeError, match="missing 1 required positional argument: 'token'"): - Agent() # Missing required token parameter + Agent() # Missing required token parameter # type: ignore[call-arg] def test_agent_init_with_none_token(self): """Test Agent initialization with None token.""" diff --git a/tests/unit/codegen/agents/test_usage_demo.py b/tests/unit/codegen/agents/test_usage_demo.py index e04bb2803..1c75fe854 100644 --- a/tests/unit/codegen/agents/test_usage_demo.py +++ b/tests/unit/codegen/agents/test_usage_demo.py @@ -39,7 +39,7 @@ def test_error_handling_no_token(): with pytest.raises(TypeError, match="missing 1 required positional argument: 'token'"): from codegen.agents.agent import Agent - Agent() # No token provided + Agent() # No token provided # type: ignore[call-arg] def test_basic_initialization_variations():