180 lines
7.1 KiB
Python
180 lines
7.1 KiB
Python
"""API Caller Module"""
|
||
import requests
|
||
from typing import Any, Dict, Optional, Union, List
|
||
|
||
from pydantic import BaseModel, Field, HttpUrl
|
||
|
||
# It's a good practice to define input/output models,
|
||
# even for internal components, using Pydantic.
|
||
|
||
class APIRequest(BaseModel):
|
||
method: str # GET, POST, PUT, DELETE, etc.
|
||
url: HttpUrl
|
||
headers: Optional[Dict[str, str]] = None
|
||
params: Optional[Dict[str, Any]] = None
|
||
json_data: Optional[Any] = None # For POST/PUT with JSON body (can be dict, list, string, etc.)
|
||
body: Optional[Any] = Field(default=None, description="Alias for json_data") # 添加别名，方便调用
|
||
data: Optional[Any] = None # For form data etc.
|
||
timeout: int = 30 # seconds
|
||
|
||
def model_post_init(self, __context):
|
||
"""初始化后处理，将 body 赋值给 json_data（如果提供了body）"""
|
||
if self.body is not None and self.json_data is None:
|
||
self.json_data = self.body
|
||
|
||
class APIResponse(BaseModel):
|
||
status_code: int
|
||
headers: Dict[str, str]
|
||
content: bytes # Raw content
|
||
json_content: Optional[Any] = None # Parsed JSON content if applicable
|
||
elapsed_time: float # in seconds
|
||
|
||
class APICaller:
|
||
"""
|
||
Responsible for executing HTTP/S API calls to the DDMS services.
|
||
"""
|
||
|
||
def __init__(self, default_timeout: int = 30, default_headers: Optional[Dict[str, str]] = None):
|
||
self.default_timeout = default_timeout
|
||
self.default_headers = default_headers or {}
|
||
|
||
def call_api(self, request_data: APIRequest) -> APIResponse:
|
||
"""
|
||
Makes an API call based on the provided request data.
|
||
|
||
Args:
|
||
request_data: An APIRequest Pydantic model instance.
|
||
|
||
Returns:
|
||
An APIResponse Pydantic model instance.
|
||
"""
|
||
merged_headers = {**self.default_headers, **(request_data.headers or {})}
|
||
timeout = request_data.timeout or self.default_timeout
|
||
|
||
try:
|
||
# 如果提供了 body，使用它作为 json 参数
|
||
json_payload = request_data.json_data
|
||
|
||
response = requests.request(
|
||
method=request_data.method.upper(),
|
||
url=str(request_data.url),
|
||
headers=merged_headers,
|
||
params=request_data.params,
|
||
json=json_payload,
|
||
data=request_data.data,
|
||
timeout=timeout
|
||
)
|
||
|
||
# 不立即引发异常，而是捕获状态码
|
||
status_code = response.status_code
|
||
|
||
json_content = None
|
||
try:
|
||
if response.headers.get('Content-Type', '').startswith('application/json'):
|
||
json_content = response.json()
|
||
except requests.exceptions.JSONDecodeError:
|
||
# Not a JSON response or invalid JSON, that's fine for some cases.
|
||
pass
|
||
|
||
return APIResponse(
|
||
status_code=status_code,
|
||
headers=dict(response.headers),
|
||
content=response.content,
|
||
json_content=json_content,
|
||
elapsed_time=response.elapsed.total_seconds()
|
||
)
|
||
except requests.exceptions.HTTPError as e:
|
||
# 处理 HTTP 错误
|
||
print(f"API call to {request_data.url} failed: {e}")
|
||
return APIResponse(
|
||
status_code=e.response.status_code,
|
||
headers=dict(e.response.headers),
|
||
content=e.response.content,
|
||
json_content=None,
|
||
elapsed_time=e.response.elapsed.total_seconds()
|
||
)
|
||
except requests.exceptions.RequestException as e:
|
||
# 处理其他请求异常
|
||
print(f"API call to {request_data.url} failed: {e}")
|
||
return APIResponse(
|
||
status_code=getattr(e.response, 'status_code', 500),
|
||
headers=dict(getattr(e.response, 'headers', {})),
|
||
content=str(e).encode(),
|
||
json_content=None,
|
||
elapsed_time=0
|
||
)
|
||
|
||
# Example Usage (can be moved to tests or main application logic)
|
||
if __name__ == '__main__':
|
||
caller = APICaller(default_headers={"X-App-Name": "DDMSComplianceSuite"})
|
||
|
||
# Example GET request
|
||
get_req_data = APIRequest(
|
||
method="GET",
|
||
url=HttpUrl("https://jsonplaceholder.typicode.com/todos/1"),
|
||
headers={"X-Request-ID": "12345"}
|
||
)
|
||
response = caller.call_api(get_req_data)
|
||
print("GET Response:")
|
||
if response.json_content:
|
||
print(f"Status: {response.status_code}, Data: {response.json_content}")
|
||
else:
|
||
print(f"Status: {response.status_code}, Content: {response.content.decode()}")
|
||
print(f"Time taken: {response.elapsed_time:.4f}s")
|
||
|
||
print("\n")
|
||
|
||
# Example POST request with json_data
|
||
post_req_data = APIRequest(
|
||
method="POST",
|
||
url=HttpUrl("https://jsonplaceholder.typicode.com/posts"),
|
||
json_data={"title": "foo", "body": "bar", "userId": 1},
|
||
headers={"Content-Type": "application/json; charset=UTF-8"}
|
||
)
|
||
response = caller.call_api(post_req_data)
|
||
print("POST Response with json_data:")
|
||
if response.json_content:
|
||
print(f"Status: {response.status_code}, Data: {response.json_content}")
|
||
else:
|
||
print(f"Status: {response.status_code}, Content: {response.content.decode()}")
|
||
print(f"Time taken: {response.elapsed_time:.4f}s")
|
||
|
||
# Example POST request with body (alias for json_data)
|
||
post_req_data_with_body = APIRequest(
|
||
method="POST",
|
||
url=HttpUrl("https://jsonplaceholder.typicode.com/posts"),
|
||
body={"title": "using body", "body": "testing body alias", "userId": 2},
|
||
headers={"Content-Type": "application/json; charset=UTF-8"}
|
||
)
|
||
response = caller.call_api(post_req_data_with_body)
|
||
print("\nPOST Response with body:")
|
||
if response.json_content:
|
||
print(f"Status: {response.status_code}, Data: {response.json_content}")
|
||
else:
|
||
print(f"Status: {response.status_code}, Content: {response.content.decode()}")
|
||
print(f"Time taken: {response.elapsed_time:.4f}s")
|
||
|
||
# Example list body request (demonstrating array body support)
|
||
array_req_data = APIRequest(
|
||
method="POST",
|
||
url=HttpUrl("https://jsonplaceholder.typicode.com/posts"),
|
||
body=["test_string", "another_value"],
|
||
headers={"Content-Type": "application/json; charset=UTF-8"}
|
||
)
|
||
response = caller.call_api(array_req_data)
|
||
print("\nPOST Response with array body:")
|
||
if response.json_content:
|
||
print(f"Status: {response.status_code}, Data: {response.json_content}")
|
||
else:
|
||
print(f"Status: {response.status_code}, Content: {response.content.decode()}")
|
||
print(f"Time taken: {response.elapsed_time:.4f}s")
|
||
|
||
# Example Error request (non-existent domain)
|
||
error_req_data = APIRequest(
|
||
method="GET",
|
||
url=HttpUrl("https://nonexistentdomain.invalid"),
|
||
)
|
||
response = caller.call_api(error_req_data)
|
||
print("\nError GET Response:")
|
||
print(f"Status: {response.status_code}, Content: {response.content.decode()}")
|
||
print(f"Time taken: {response.elapsed_time:.4f}s") |