125 lines
4.1 KiB
Python
125 lines
4.1 KiB
Python
|
import httpx
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
|
||
|
|
class RecordingAsyncClient:
|
||
|
|
def __init__(self, save_dir: str):
|
||
|
|
self.save_dir = save_dir
|
||
|
|
os.makedirs(self.save_dir, exist_ok=True)
|
||
|
|
self.client = None # Will be initialized in __aenter__
|
||
|
|
|
||
|
|
async def __aenter__(self):
|
||
|
|
# Initialize the actual AsyncClient when entering the context manager
|
||
|
|
self.client = httpx.AsyncClient()
|
||
|
|
return self
|
||
|
|
|
||
|
|
async def __aexit__(self, exc_type, exc_value, traceback):
|
||
|
|
# Ensure the client is closed when exiting the context manager
|
||
|
|
await self.client.aclose()
|
||
|
|
|
||
|
|
async def request(self, method: str, url: str, **kwargs):
|
||
|
|
# Send the actual request
|
||
|
|
response = await self.client.request(method, url, **kwargs)
|
||
|
|
|
||
|
|
# Record the request and response
|
||
|
|
record = {
|
||
|
|
"request": {
|
||
|
|
"method": method,
|
||
|
|
"url": url,
|
||
|
|
"headers": dict(response.request.headers),
|
||
|
|
"content": response.request.content.decode('utf-8', errors='ignore'),
|
||
|
|
},
|
||
|
|
"response": {
|
||
|
|
"status_code": response.status_code,
|
||
|
|
"headers": dict(response.headers),
|
||
|
|
"content": response.text,
|
||
|
|
"cookies": dict(response.cookies),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
# Generate a filename based on the URL and method
|
||
|
|
record_file = os.path.join(self.save_dir, f"{method}_{url.replace('/', '_')}.json")
|
||
|
|
|
||
|
|
# Save the record to a file
|
||
|
|
with open(record_file, 'w') as f:
|
||
|
|
json.dump(record, f, indent=4)
|
||
|
|
|
||
|
|
return response
|
||
|
|
|
||
|
|
async def get(self, url: str, **kwargs):
|
||
|
|
return await self.request("GET", url, **kwargs)
|
||
|
|
|
||
|
|
async def post(self, url: str, **kwargs):
|
||
|
|
return await self.request("POST", url, **kwargs)
|
||
|
|
|
||
|
|
async def put(self, url: str, **kwargs):
|
||
|
|
return await self.request("PUT", url, **kwargs)
|
||
|
|
|
||
|
|
async def delete(self, url: str, **kwargs):
|
||
|
|
return await self.request("DELETE", url, **kwargs)
|
||
|
|
|
||
|
|
from unittest.mock import Mock
|
||
|
|
import os
|
||
|
|
import json
|
||
|
|
|
||
|
|
class MockAsyncClient:
|
||
|
|
def __init__(self, load_dir: str):
|
||
|
|
self.load_dir = load_dir
|
||
|
|
|
||
|
|
async def __aenter__(self):
|
||
|
|
# No actual client to initialize, just return the instance
|
||
|
|
return self
|
||
|
|
|
||
|
|
async def __aexit__(self, exc_type, exc_value, traceback):
|
||
|
|
# No actual client to close
|
||
|
|
pass
|
||
|
|
|
||
|
|
async def request(self, method: str, url: str, **kwargs):
|
||
|
|
# Generate the filename based on the URL and method
|
||
|
|
record_file = os.path.join(self.load_dir, f"{method}_{url.replace('/', '_')}.json")
|
||
|
|
|
||
|
|
if not os.path.exists(record_file):
|
||
|
|
raise FileNotFoundError(f"Recorded response not found for {method} {url}")
|
||
|
|
|
||
|
|
# Load the recorded response from the file
|
||
|
|
with open(record_file, 'r') as f:
|
||
|
|
record = json.load(f)
|
||
|
|
|
||
|
|
# Create a mock response object
|
||
|
|
mock_response = Mock()
|
||
|
|
|
||
|
|
# Mock the status code
|
||
|
|
mock_response.status_code = record['response']['status_code']
|
||
|
|
|
||
|
|
# Mock the json method to return the content as a parsed JSON
|
||
|
|
def mock_json():
|
||
|
|
try:
|
||
|
|
return json.loads(record['response']['content'])
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
return record['response']['content']
|
||
|
|
|
||
|
|
mock_response.json = mock_json
|
||
|
|
|
||
|
|
# Mock the cookies as a dictionary
|
||
|
|
mock_response.cookies = record['response']['cookies']
|
||
|
|
|
||
|
|
# Mock the headers as a dictionary
|
||
|
|
mock_response.headers = record['response']['headers']
|
||
|
|
|
||
|
|
# Mock the text attribute
|
||
|
|
mock_response.text = record['response']['content']
|
||
|
|
|
||
|
|
return mock_response
|
||
|
|
|
||
|
|
async def get(self, url: str, **kwargs):
|
||
|
|
return await self.request("GET", url, **kwargs)
|
||
|
|
|
||
|
|
async def post(self, url: str, **kwargs):
|
||
|
|
return await self.request("POST", url, **kwargs)
|
||
|
|
|
||
|
|
async def put(self, url: str, **kwargs):
|
||
|
|
return await self.request("PUT", url, **kwargs)
|
||
|
|
|
||
|
|
async def delete(self, url: str, **kwargs):
|
||
|
|
return await self.request("DELETE", url, **kwargs)
|