"""
Utilities for mocking async context managers and Playwright objects.
"""
from typing import Any, Dict, Optional, AsyncIterator
from unittest.mock import AsyncMock, MagicMock
class AsyncContextManagerMock:
"""Mock for async context managers."""
def __init__(self, mock_obj: Any = None):
self.mock_obj = mock_obj if mock_obj is not None else AsyncMock()
async def __aenter__(self) -> Any:
return self.mock_obj
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
pass
class PlaywrightMock:
"""Mock for Playwright."""
def __init__(self):
self.chromium = AsyncMock()
self.chromium.launch = AsyncMock()
self.stop = AsyncMock()
@staticmethod
async def start() -> 'PlaywrightMock':
"""Return the mock instance."""
return PlaywrightMock()
def create_mock_response(status: int = 200, text: str = "", headers: Optional[Dict[str, str]] = None) -> AsyncMock:
"""Create a mock aiohttp response."""
mock_response = AsyncMock()
mock_response.status = status
mock_response.reason = "OK" if 200 <= status < 300 else "Error"
mock_response.text = AsyncMock(return_value=text)
mock_response.headers = headers or {"content-type": "text/html"}
return mock_response
class MockClientSession:
"""Mock for aiohttp.ClientSession with proper async context manager support."""
def __init__(self):
self.closed = False
self.get_responses = {}
self.post_responses = {}
def set_response(self, url: str, response: AsyncMock, method: str = "get"):
"""Set a response for a specific URL."""
if method == "get":
self.get_responses[url] = response
else:
self.post_responses[url] = response
def get(self, url: str, **kwargs):
"""Mock get method."""
if url in self.get_responses:
context_mock = AsyncContextManagerMock(self.get_responses[url])
else:
# Default response
context_mock = AsyncContextManagerMock(create_mock_response(200, "<html><body>Default response</body></html>"))
return context_mock
def post(self, url: str, **kwargs):
"""Mock post method."""
if url in self.post_responses:
context_mock = AsyncContextManagerMock(self.post_responses[url])
else:
# Default response
context_mock = AsyncContextManagerMock(create_mock_response(200, '{"status": "success"}'))
return context_mock
async def close(self) -> None:
"""Close the session."""
self.closed = True
