conftest.py•4.84 kB
#!/usr/bin/env python3
"""
Pytest configuration and shared fixtures for MetasploitMCP tests.
"""
import pytest
import sys
import os
from unittest.mock import Mock, patch
# Add the project root to Python path
sys.path.insert(0, os.path.dirname(__file__))
def pytest_configure(config):
"""Configure pytest with custom settings."""
# Mock external dependencies that might not be available
mock_modules = [
'uvicorn',
'fastapi',
'mcp.server.fastmcp',
'mcp.server.sse',
'pymetasploit3.msfrpc',
'starlette.applications',
'starlette.routing',
'mcp.server.session'
]
for module in mock_modules:
if module not in sys.modules:
sys.modules[module] = Mock()
def pytest_collection_modifyitems(config, items):
"""Modify test collection to add markers automatically."""
for item in items:
# Add unit marker to test_options_parsing and test_helpers
if "test_options_parsing" in item.nodeid or "test_helpers" in item.nodeid:
item.add_marker(pytest.mark.unit)
# Add integration marker to integration tests
if "test_tools_integration" in item.nodeid:
item.add_marker(pytest.mark.integration)
# Mark network tests
if "network" in item.name.lower():
item.add_marker(pytest.mark.network)
# Mark slow tests
if any(keyword in item.name.lower() for keyword in ["slow", "timeout", "long"]):
item.add_marker(pytest.mark.slow)
@pytest.fixture(scope="session")
def mock_msf_environment():
"""Session-scoped fixture that provides a complete mock MSF environment."""
class MockMsfRpcClient:
def __init__(self):
self.modules = Mock()
self.core = Mock()
self.sessions = Mock()
self.jobs = Mock()
self.consoles = Mock()
# Default return values
self.core.version = {'version': '6.3.0'}
self.modules.exploits = []
self.modules.payloads = []
self.sessions.list.return_value = {}
self.jobs.list.return_value = {}
class MockMsfConsole:
def __init__(self, cid='test-console'):
self.cid = cid
def read(self):
return {'data': '', 'prompt': 'msf6 > ', 'busy': False}
def write(self, command):
return True
class MockMsfRpcError(Exception):
pass
# Apply mocks
with patch.dict('sys.modules', {
'pymetasploit3.msfrpc': Mock(
MsfRpcClient=MockMsfRpcClient,
MsfConsole=MockMsfConsole,
MsfRpcError=MockMsfRpcError
)
}):
yield {
'client_class': MockMsfRpcClient,
'console_class': MockMsfConsole,
'error_class': MockMsfRpcError
}
@pytest.fixture
def mock_logger():
"""Fixture providing a mock logger."""
with patch('MetasploitMCP.logger') as mock_log:
yield mock_log
@pytest.fixture
def temp_payload_dir(tmp_path):
"""Fixture providing a temporary directory for payload saves."""
payload_dir = tmp_path / "payloads"
payload_dir.mkdir()
with patch('MetasploitMCP.PAYLOAD_SAVE_DIR', str(payload_dir)):
yield str(payload_dir)
@pytest.fixture
def mock_asyncio_to_thread():
"""Fixture to mock asyncio.to_thread for testing."""
async def mock_to_thread(func, *args, **kwargs):
return func(*args, **kwargs)
with patch('asyncio.to_thread', side_effect=mock_to_thread):
yield
@pytest.fixture
def capture_logs(caplog):
"""Fixture to capture and provide log output."""
import logging
caplog.set_level(logging.DEBUG)
return caplog
# Command line options
def pytest_addoption(parser):
"""Add custom command line options."""
parser.addoption(
"--run-slow",
action="store_true",
default=False,
help="Run slow tests"
)
parser.addoption(
"--run-network",
action="store_true",
default=False,
help="Run tests that require network access"
)
def pytest_runtest_setup(item):
"""Setup hook to skip tests based on markers and options."""
if "slow" in item.keywords and not item.config.getoption("--run-slow"):
pytest.skip("Skipping slow test (use --run-slow to run)")
if "network" in item.keywords and not item.config.getoption("--run-network"):
pytest.skip("Skipping network test (use --run-network to run)")
# Test environment setup
@pytest.fixture(autouse=True)
def reset_msf_client():
"""Automatically reset the global MSF client between tests."""
with patch('MetasploitMCP._msf_client_instance', None):
yield