21 changed files with 5769 additions and 38 deletions
@ -0,0 +1,112 @@ |
|||||
|
#!/usr/bin/env python3 |
||||
|
""" |
||||
|
Hook to capture session start/end events. |
||||
|
Cross-platform support for Windows, macOS, and Linux. |
||||
|
""" |
||||
|
import json |
||||
|
import sys |
||||
|
import os |
||||
|
import subprocess |
||||
|
from datetime import datetime, timezone |
||||
|
from claude_code_capture_utils import get_log_file_path, add_ab_metadata |
||||
|
|
||||
|
def get_git_metadata(repo_dir): |
||||
|
"""Get current git commit and branch.""" |
||||
|
try: |
||||
|
# Get current commit hash |
||||
|
commit_result = subprocess.run( |
||||
|
['git', 'rev-parse', 'HEAD'], |
||||
|
cwd=repo_dir, |
||||
|
capture_output=True, |
||||
|
text=True, |
||||
|
timeout=10 |
||||
|
) |
||||
|
|
||||
|
# Get current branch |
||||
|
branch_result = subprocess.run( |
||||
|
['git', 'branch', '--show-current'], |
||||
|
cwd=repo_dir, |
||||
|
capture_output=True, |
||||
|
text=True, |
||||
|
timeout=10 |
||||
|
) |
||||
|
|
||||
|
git_metadata = { |
||||
|
"base_commit": commit_result.stdout.strip() if commit_result.returncode == 0 else None, |
||||
|
"branch": branch_result.stdout.strip() if branch_result.returncode == 0 else None, |
||||
|
"timestamp": datetime.now(timezone.utc).isoformat() |
||||
|
} |
||||
|
|
||||
|
if git_metadata["base_commit"]: |
||||
|
return git_metadata |
||||
|
else: |
||||
|
return None |
||||
|
|
||||
|
except Exception as e: |
||||
|
print(f"Warning: Could not capture git metadata: {e}", file=sys.stderr) |
||||
|
return None |
||||
|
|
||||
|
def main(): |
||||
|
try: |
||||
|
if len(sys.argv) < 2: |
||||
|
print("Usage: capture_session_event.py [start|end]", file=sys.stderr) |
||||
|
sys.exit(1) |
||||
|
|
||||
|
event_type = sys.argv[1].lower() |
||||
|
if event_type not in ["start", "end"]: |
||||
|
print("Event type must be 'start' or 'end'", file=sys.stderr) |
||||
|
sys.exit(1) |
||||
|
|
||||
|
input_data = json.load(sys.stdin) |
||||
|
|
||||
|
session_id = input_data.get("session_id", "unknown") |
||||
|
transcript_path = input_data.get("transcript_path", "") |
||||
|
cwd = input_data.get("cwd", "") |
||||
|
|
||||
|
if event_type == "start": |
||||
|
# Session start: capture git metadata |
||||
|
git_metadata = get_git_metadata(cwd) |
||||
|
|
||||
|
log_entry = { |
||||
|
"type": "session_start", |
||||
|
"timestamp": datetime.now(timezone.utc).isoformat(), |
||||
|
"session_id": session_id, |
||||
|
"transcript_path": transcript_path, |
||||
|
"cwd": cwd, |
||||
|
"git_metadata": git_metadata |
||||
|
} |
||||
|
|
||||
|
log_entry = add_ab_metadata(log_entry, cwd) |
||||
|
|
||||
|
if git_metadata: |
||||
|
print(f"[OK] Captured git metadata: {git_metadata['base_commit'][:8]} on {git_metadata['branch']}") |
||||
|
|
||||
|
# Write session_start event |
||||
|
log_file = get_log_file_path(session_id, cwd) |
||||
|
with open(log_file, "a", encoding="utf-8") as f: |
||||
|
f.write(json.dumps(log_entry) + "\n") |
||||
|
|
||||
|
elif event_type == "end": |
||||
|
# Session end: log the event |
||||
|
log_entry = { |
||||
|
"type": "session_end", |
||||
|
"timestamp": datetime.now(timezone.utc).isoformat(), |
||||
|
"session_id": session_id, |
||||
|
"transcript_path": transcript_path, |
||||
|
"cwd": cwd, |
||||
|
"reason": input_data.get("reason", "") |
||||
|
} |
||||
|
|
||||
|
log_entry = add_ab_metadata(log_entry, cwd) |
||||
|
|
||||
|
# Write session_end event |
||||
|
log_file = get_log_file_path(session_id, cwd) |
||||
|
with open(log_file, "a", encoding="utf-8") as f: |
||||
|
f.write(json.dumps(log_entry) + "\n") |
||||
|
|
||||
|
except Exception as e: |
||||
|
print(f"[ERROR] Session {event_type}: {e}", file=sys.stderr) |
||||
|
sys.exit(1) |
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
main() |
||||
@ -0,0 +1,93 @@ |
|||||
|
#!/usr/bin/env python3 |
||||
|
""" |
||||
|
Utility functions for A/B testing hooks. |
||||
|
Cross-platform support for Windows, macOS, and Linux. |
||||
|
""" |
||||
|
import json |
||||
|
import os |
||||
|
from pathlib import Path |
||||
|
|
||||
|
def detect_model_lane(cwd): |
||||
|
"""Detect if we're in model_a or model_b directory.""" |
||||
|
path_parts = Path(cwd).parts |
||||
|
if 'model_a' in path_parts: |
||||
|
return 'model_a' |
||||
|
elif 'model_b' in path_parts: |
||||
|
return 'model_b' |
||||
|
return None |
||||
|
|
||||
|
def get_experiment_root(cwd): |
||||
|
"""Get the experiment root directory (parent of model_a/model_b).""" |
||||
|
current_path = Path(cwd) |
||||
|
|
||||
|
# Check if we're inside a model_a or model_b directory |
||||
|
# Look for the parent that contains both model_a and model_b |
||||
|
for parent in [current_path] + list(current_path.parents): |
||||
|
if (parent / 'model_a').exists() and (parent / 'model_b').exists(): |
||||
|
return str(parent) |
||||
|
# Also check one level up (in case we're inside the cloned repo) |
||||
|
parent_up = parent.parent |
||||
|
if (parent_up / 'model_a').exists() and (parent_up / 'model_b').exists(): |
||||
|
return str(parent_up) |
||||
|
return None |
||||
|
|
||||
|
def read_manifest(experiment_root): |
||||
|
"""Read the manifest.json file to get task_id and model assignments.""" |
||||
|
try: |
||||
|
manifest_path = os.path.join(experiment_root, 'manifest.json') |
||||
|
if os.path.exists(manifest_path): |
||||
|
with open(manifest_path, 'r', encoding='utf-8') as f: |
||||
|
return json.load(f) |
||||
|
except Exception: |
||||
|
pass |
||||
|
return {} |
||||
|
|
||||
|
def get_ab_metadata(cwd): |
||||
|
"""Get A/B testing metadata (task_id, model_lane, model_name) from current directory.""" |
||||
|
model_lane = detect_model_lane(cwd) |
||||
|
experiment_root = get_experiment_root(cwd) |
||||
|
|
||||
|
if not model_lane or not experiment_root: |
||||
|
return {} |
||||
|
|
||||
|
manifest = read_manifest(experiment_root) |
||||
|
|
||||
|
metadata = { |
||||
|
"task_id": manifest.get("task_id"), |
||||
|
"model_lane": model_lane, |
||||
|
"experiment_root": experiment_root |
||||
|
} |
||||
|
|
||||
|
# Get model name from assignments |
||||
|
assignments = manifest.get("assignments", {}) |
||||
|
if model_lane in assignments: |
||||
|
metadata["model_name"] = assignments[model_lane] |
||||
|
|
||||
|
return metadata |
||||
|
|
||||
|
def get_log_file_path(session_id, cwd): |
||||
|
"""Get the correct log file path for A/B testing (routes to model-specific directory).""" |
||||
|
model_lane = detect_model_lane(cwd) |
||||
|
experiment_root = get_experiment_root(cwd) |
||||
|
|
||||
|
if model_lane and experiment_root: |
||||
|
# Route to model-specific logs directory |
||||
|
logs_dir = os.path.join(experiment_root, "logs", model_lane) |
||||
|
os.makedirs(logs_dir, exist_ok=True) |
||||
|
return os.path.join(logs_dir, f"session_{session_id}.jsonl") |
||||
|
else: |
||||
|
# Fallback to current behavior |
||||
|
project_dir = os.environ.get("CLAUDE_PROJECT_DIR", os.getcwd()) |
||||
|
logs_dir = os.path.join(project_dir, "logs") |
||||
|
os.makedirs(logs_dir, exist_ok=True) |
||||
|
return os.path.join(logs_dir, f"session_{session_id}.jsonl") |
||||
|
|
||||
|
def add_ab_metadata(event, cwd): |
||||
|
"""Add A/B testing metadata to an event.""" |
||||
|
ab_metadata = get_ab_metadata(cwd) |
||||
|
if ab_metadata: |
||||
|
for key, value in ab_metadata.items(): |
||||
|
if value is not None: |
||||
|
event[key] = value |
||||
|
return event |
||||
|
|
||||
@ -0,0 +1,720 @@ |
|||||
|
#!/usr/bin/env python3 |
||||
|
""" |
||||
|
Process raw transcript to extract all messages and generate summaries. |
||||
|
Handles both incremental (Stop event) and final (SessionEnd) processing. |
||||
|
Cross-platform support for Windows, macOS, and Linux. |
||||
|
""" |
||||
|
import json |
||||
|
import sys |
||||
|
import os |
||||
|
import shutil |
||||
|
import subprocess |
||||
|
from datetime import datetime, timezone |
||||
|
from pathlib import Path |
||||
|
from collections import defaultdict |
||||
|
from claude_code_capture_utils import get_log_file_path, add_ab_metadata, detect_model_lane, get_experiment_root |
||||
|
|
||||
|
def read_and_process_raw_transcript(transcript_path): |
||||
|
""" |
||||
|
Read raw transcript and extract all unique messages. |
||||
|
Returns deduplicated messages with last occurrence (final state). |
||||
|
Also extracts thinking blocks as separate entries. |
||||
|
""" |
||||
|
if not os.path.exists(transcript_path): |
||||
|
return [] |
||||
|
|
||||
|
# Track assistant messages by message.id (they have IDs, can have duplicates) |
||||
|
assistant_messages = {} |
||||
|
# Track thinking blocks separately (won't be counted in token usage) |
||||
|
thinking_blocks = {} |
||||
|
# Track user messages by uuid (they don't have message.id, use uuid) |
||||
|
user_messages = {} |
||||
|
|
||||
|
try: |
||||
|
with open(transcript_path, 'r', encoding='utf-8') as f: |
||||
|
for line in f: |
||||
|
line = line.strip() |
||||
|
if not line: |
||||
|
continue |
||||
|
|
||||
|
try: |
||||
|
event = json.loads(line) |
||||
|
event_type = event.get('type') |
||||
|
message = event.get('message', {}) |
||||
|
|
||||
|
# Process assistant messages (have message.id) |
||||
|
if event_type == 'assistant': |
||||
|
msg_id = message.get('id') |
||||
|
if msg_id: |
||||
|
# Check for thinking blocks in content |
||||
|
content = message.get('content', []) |
||||
|
if isinstance(content, list): |
||||
|
for item in content: |
||||
|
if isinstance(item, dict) and item.get('type') == 'thinking': |
||||
|
# Extract thinking block as separate entry |
||||
|
thinking_entry = { |
||||
|
'type': 'assistant_thinking', |
||||
|
'timestamp': event.get('timestamp'), |
||||
|
'message_id': msg_id, |
||||
|
'thinking_content': item.get('thinking', ''), |
||||
|
'session_id': event.get('sessionId'), |
||||
|
'cwd': event.get('cwd') |
||||
|
} |
||||
|
# Use message_id as key (one thinking block per message) |
||||
|
thinking_blocks[msg_id] = thinking_entry |
||||
|
|
||||
|
# Store/overwrite with last occurrence (streaming) |
||||
|
assistant_msg = { |
||||
|
'type': event_type, |
||||
|
'timestamp': event.get('timestamp'), |
||||
|
'message': message, |
||||
|
'session_id': event.get('sessionId'), |
||||
|
'cwd': event.get('cwd') |
||||
|
} |
||||
|
|
||||
|
# Preserve stop_reason from message if present |
||||
|
if message.get('stop_reason'): |
||||
|
assistant_msg['stop_reason'] = message['stop_reason'] |
||||
|
|
||||
|
assistant_messages[msg_id] = assistant_msg |
||||
|
|
||||
|
# Process user messages (use uuid as key, no message.id) |
||||
|
elif event_type == 'user': |
||||
|
uuid = event.get('uuid') |
||||
|
if uuid: |
||||
|
user_msg = { |
||||
|
'type': event_type, |
||||
|
'timestamp': event.get('timestamp'), |
||||
|
'message': message, |
||||
|
'session_id': event.get('sessionId'), |
||||
|
'cwd': event.get('cwd'), |
||||
|
'uuid': uuid |
||||
|
} |
||||
|
|
||||
|
# Preserve thinking metadata if present |
||||
|
if 'thinkingMetadata' in event: |
||||
|
user_msg['thinkingMetadata'] = event['thinkingMetadata'] |
||||
|
|
||||
|
# Preserve isMeta flag if present |
||||
|
if event.get('isMeta'): |
||||
|
user_msg['isMeta'] = event['isMeta'] |
||||
|
|
||||
|
user_messages[uuid] = user_msg |
||||
|
|
||||
|
except json.JSONDecodeError: |
||||
|
continue |
||||
|
|
||||
|
except Exception as e: |
||||
|
print(f"[ERROR] Reading raw transcript: {e}", file=sys.stderr) |
||||
|
return [] |
||||
|
|
||||
|
# Combine all: assistant messages + thinking blocks + user messages |
||||
|
# Thinking blocks inserted right before their corresponding assistant message |
||||
|
all_messages = [] |
||||
|
|
||||
|
# First, add all messages with their thinking blocks properly ordered |
||||
|
assistant_list = list(assistant_messages.values()) |
||||
|
user_list = list(user_messages.values()) |
||||
|
|
||||
|
# Combine and sort all by timestamp |
||||
|
combined = assistant_list + user_list |
||||
|
combined.sort(key=lambda m: m.get('timestamp', '')) |
||||
|
|
||||
|
# Insert thinking blocks right before their parent assistant message |
||||
|
for msg in combined: |
||||
|
if msg['type'] == 'assistant': |
||||
|
msg_id = msg['message'].get('id') |
||||
|
# If there's a thinking block for this message, insert it first |
||||
|
if msg_id in thinking_blocks: |
||||
|
all_messages.append(thinking_blocks[msg_id]) |
||||
|
all_messages.append(msg) |
||||
|
|
||||
|
return all_messages |
||||
|
|
||||
|
def aggregate_token_usage(messages): |
||||
|
"""Aggregate token usage from all assistant messages. |
||||
|
Note: assistant_thinking entries are explicitly excluded to avoid double-counting. |
||||
|
""" |
||||
|
total_usage = { |
||||
|
'total_input_tokens': 0, |
||||
|
'total_output_tokens': 0, |
||||
|
'total_cache_creation_tokens': 0, |
||||
|
'total_cache_read_tokens': 0, |
||||
|
'total_ephemeral_5m_tokens': 0, |
||||
|
'total_ephemeral_1h_tokens': 0, |
||||
|
'service_tier': None |
||||
|
} |
||||
|
|
||||
|
for msg_data in messages: |
||||
|
# Only count tokens from 'assistant' type, NOT 'assistant_thinking' |
||||
|
# Thinking tokens are already included in the parent assistant message's output_tokens |
||||
|
if msg_data['type'] == 'assistant': |
||||
|
message = msg_data['message'] |
||||
|
usage = message.get('usage', {}) |
||||
|
|
||||
|
if usage: |
||||
|
total_usage['total_input_tokens'] += usage.get('input_tokens', 0) |
||||
|
total_usage['total_output_tokens'] += usage.get('output_tokens', 0) |
||||
|
total_usage['total_cache_creation_tokens'] += usage.get('cache_creation_input_tokens', 0) |
||||
|
total_usage['total_cache_read_tokens'] += usage.get('cache_read_input_tokens', 0) |
||||
|
|
||||
|
cache_creation = usage.get('cache_creation', {}) |
||||
|
total_usage['total_ephemeral_5m_tokens'] += cache_creation.get('ephemeral_5m_input_tokens', 0) |
||||
|
total_usage['total_ephemeral_1h_tokens'] += cache_creation.get('ephemeral_1h_input_tokens', 0) |
||||
|
|
||||
|
if usage.get('service_tier'): |
||||
|
total_usage['service_tier'] = usage.get('service_tier') |
||||
|
|
||||
|
# Add calculated total |
||||
|
total_usage['total_actual_input_tokens'] = ( |
||||
|
total_usage['total_input_tokens'] + |
||||
|
total_usage['total_cache_creation_tokens'] + |
||||
|
total_usage['total_cache_read_tokens'] |
||||
|
) |
||||
|
|
||||
|
return total_usage |
||||
|
|
||||
|
def analyze_tool_calls(messages): |
||||
|
"""Extract tool call metrics from messages.""" |
||||
|
tool_calls = defaultdict(int) |
||||
|
tool_results = defaultdict(int) |
||||
|
|
||||
|
for msg_data in messages: |
||||
|
# Skip entries without 'message' key (e.g., assistant_thinking) |
||||
|
if 'message' not in msg_data: |
||||
|
continue |
||||
|
message = msg_data['message'] |
||||
|
content = message.get('content', []) |
||||
|
|
||||
|
if not isinstance(content, list): |
||||
|
continue |
||||
|
|
||||
|
for item in content: |
||||
|
if not isinstance(item, dict): |
||||
|
continue |
||||
|
|
||||
|
if item.get('type') == 'tool_use': |
||||
|
tool_name = item.get('name', 'unknown') |
||||
|
tool_calls[tool_name] += 1 |
||||
|
|
||||
|
elif item.get('type') == 'tool_result': |
||||
|
# Try to infer tool name from context (simplified) |
||||
|
tool_results['total'] += 1 |
||||
|
|
||||
|
return { |
||||
|
'tool_calls_by_type': dict(tool_calls), |
||||
|
'total_tool_calls': sum(tool_calls.values()), |
||||
|
'total_tool_results': tool_results.get('total', 0) |
||||
|
} |
||||
|
|
||||
|
def analyze_thinking_usage(messages, transcript_path): |
||||
|
"""Analyze thinking mode usage in messages.""" |
||||
|
thinking_stats = { |
||||
|
'thinking_enabled_turns': 0, |
||||
|
'thinking_disabled_turns': 0, |
||||
|
'assistant_with_thinking_blocks': 0, |
||||
|
'thinking_levels': defaultdict(int) |
||||
|
} |
||||
|
|
||||
|
# Track which turns had thinking enabled (from user thinkingMetadata) |
||||
|
for msg_data in messages: |
||||
|
if msg_data['type'] == 'user' and 'thinkingMetadata' in msg_data: |
||||
|
metadata = msg_data['thinkingMetadata'] |
||||
|
if not metadata.get('disabled', True): |
||||
|
thinking_stats['thinking_enabled_turns'] += 1 |
||||
|
level = metadata.get('level', 'none') |
||||
|
thinking_stats['thinking_levels'][level] += 1 |
||||
|
else: |
||||
|
thinking_stats['thinking_disabled_turns'] += 1 |
||||
|
|
||||
|
# Count assistant messages with thinking blocks |
||||
|
# Check ALL occurrences in raw transcript (not just final deduplicated state) |
||||
|
assistant_msg_ids_with_thinking = set() |
||||
|
|
||||
|
try: |
||||
|
if os.path.exists(transcript_path): |
||||
|
with open(transcript_path, 'r', encoding='utf-8') as f: |
||||
|
for line in f: |
||||
|
line = line.strip() |
||||
|
if not line: |
||||
|
continue |
||||
|
try: |
||||
|
event = json.loads(line) |
||||
|
if event.get('type') == 'assistant': |
||||
|
message = event.get('message', {}) |
||||
|
msg_id = message.get('id') |
||||
|
content = message.get('content', []) |
||||
|
|
||||
|
if msg_id and isinstance(content, list): |
||||
|
# Check if this occurrence has thinking |
||||
|
has_thinking = any( |
||||
|
isinstance(item, dict) and item.get('type') == 'thinking' |
||||
|
for item in content |
||||
|
) |
||||
|
if has_thinking: |
||||
|
assistant_msg_ids_with_thinking.add(msg_id) |
||||
|
except: |
||||
|
continue |
||||
|
except Exception: |
||||
|
pass |
||||
|
|
||||
|
thinking_stats['assistant_with_thinking_blocks'] = len(assistant_msg_ids_with_thinking) |
||||
|
|
||||
|
return { |
||||
|
'thinking_enabled_turns': thinking_stats['thinking_enabled_turns'], |
||||
|
'thinking_disabled_turns': thinking_stats['thinking_disabled_turns'], |
||||
|
'assistant_with_thinking_blocks': thinking_stats['assistant_with_thinking_blocks'], |
||||
|
'thinking_levels': dict(thinking_stats['thinking_levels']) |
||||
|
} |
||||
|
|
||||
|
def calculate_git_metrics(cwd, base_commit): |
||||
|
"""Calculate git metrics from diff.""" |
||||
|
try: |
||||
|
original_cwd = os.getcwd() |
||||
|
os.chdir(cwd) |
||||
|
|
||||
|
if not base_commit: |
||||
|
os.chdir(original_cwd) |
||||
|
return {} |
||||
|
|
||||
|
# Add untracked files |
||||
|
excluded_patterns = ['.claude/', '__pycache__/', 'node_modules/', '.mypy_cache/', |
||||
|
'.pytest_cache/', '.DS_Store', '.vscode/', '.idea/'] |
||||
|
|
||||
|
untracked_result = subprocess.run( |
||||
|
['git', 'ls-files', '--others', '--exclude-standard'], |
||||
|
capture_output=True, text=True, timeout=30 |
||||
|
) |
||||
|
|
||||
|
if untracked_result.returncode == 0 and untracked_result.stdout.strip(): |
||||
|
untracked_files = [ |
||||
|
f.strip() for f in untracked_result.stdout.strip().split('\n') |
||||
|
if f.strip() and not any(pattern in f for pattern in excluded_patterns) |
||||
|
] |
||||
|
|
||||
|
for file in untracked_files: |
||||
|
subprocess.run(['git', 'add', '-N', file], capture_output=True, timeout=5) |
||||
|
|
||||
|
# Calculate numstat |
||||
|
result = subprocess.run( |
||||
|
['git', 'diff', '--numstat', base_commit, '--', '.', |
||||
|
':!.claude', ':!**/.mypy_cache', ':!**/__pycache__', ':!**/.pytest_cache', |
||||
|
':!**/.DS_Store', ':!**/node_modules', ':!**/.vscode', ':!**/.idea'], |
||||
|
capture_output=True, text=True, timeout=30 |
||||
|
) |
||||
|
|
||||
|
os.chdir(original_cwd) |
||||
|
|
||||
|
if result.returncode != 0: |
||||
|
return {} |
||||
|
|
||||
|
lines = result.stdout.strip().split('\n') if result.stdout.strip() else [] |
||||
|
files_changed = 0 |
||||
|
total_lines_changed = 0 |
||||
|
|
||||
|
for line in lines: |
||||
|
if line.strip(): |
||||
|
parts = line.split('\t') |
||||
|
if len(parts) >= 3: |
||||
|
try: |
||||
|
added = int(parts[0]) if parts[0] != '-' else 0 |
||||
|
removed = int(parts[1]) if parts[1] != '-' else 0 |
||||
|
files_changed += 1 |
||||
|
total_lines_changed += added + removed |
||||
|
except ValueError: |
||||
|
continue |
||||
|
|
||||
|
return { |
||||
|
"files_changed_count": files_changed, |
||||
|
"lines_of_code_changed_count": total_lines_changed |
||||
|
} |
||||
|
|
||||
|
except Exception as e: |
||||
|
print(f"Warning: Could not calculate git metrics: {e}", file=sys.stderr) |
||||
|
if 'original_cwd' in locals(): |
||||
|
os.chdir(original_cwd) |
||||
|
return {} |
||||
|
|
||||
|
def copy_raw_transcript(transcript_path, session_id, cwd): |
||||
|
"""Copy raw transcript to logs folder.""" |
||||
|
try: |
||||
|
source_path = Path(transcript_path) |
||||
|
if not source_path.exists(): |
||||
|
print(f"Warning: Raw transcript not found at {source_path}", file=sys.stderr) |
||||
|
return False |
||||
|
|
||||
|
model_lane = detect_model_lane(cwd) |
||||
|
experiment_root = get_experiment_root(cwd) |
||||
|
|
||||
|
if model_lane and experiment_root: |
||||
|
logs_dir = Path(experiment_root) / "logs" / model_lane |
||||
|
logs_dir.mkdir(parents=True, exist_ok=True) |
||||
|
dest_path = logs_dir / f"session_{session_id}_raw.jsonl" |
||||
|
else: |
||||
|
project_dir = os.environ.get('CLAUDE_PROJECT_DIR', os.getcwd()) |
||||
|
logs_dir = Path(project_dir) / "logs" |
||||
|
logs_dir.mkdir(exist_ok=True) |
||||
|
dest_path = logs_dir / f"session_{session_id}_raw.jsonl" |
||||
|
|
||||
|
shutil.copy2(source_path, dest_path) |
||||
|
print(f"[OK] Copied raw transcript to {dest_path}") |
||||
|
return True |
||||
|
|
||||
|
except Exception as e: |
||||
|
print(f"[ERROR] Copying raw transcript: {e}", file=sys.stderr) |
||||
|
return False |
||||
|
|
||||
|
def get_base_commit_from_log(log_file): |
||||
|
"""Extract base commit from session_start event.""" |
||||
|
try: |
||||
|
if not os.path.exists(log_file): |
||||
|
return None |
||||
|
|
||||
|
with open(log_file, 'r', encoding='utf-8') as f: |
||||
|
for line in f: |
||||
|
line = line.strip() |
||||
|
if line: |
||||
|
try: |
||||
|
event = json.loads(line) |
||||
|
if event.get('type') == 'session_start': |
||||
|
git_metadata = event.get('git_metadata', {}) |
||||
|
return git_metadata.get('base_commit') |
||||
|
except json.JSONDecodeError: |
||||
|
continue |
||||
|
return None |
||||
|
except Exception: |
||||
|
return None |
||||
|
|
||||
|
def main(): |
||||
|
try: |
||||
|
if len(sys.argv) < 2: |
||||
|
print("Usage: process_transcript.py [incremental|final]", file=sys.stderr) |
||||
|
sys.exit(1) |
||||
|
|
||||
|
mode = sys.argv[1].lower() |
||||
|
if mode not in ["incremental", "final"]: |
||||
|
print("Mode must be 'incremental' or 'final'", file=sys.stderr) |
||||
|
sys.exit(1) |
||||
|
|
||||
|
input_data = json.load(sys.stdin) |
||||
|
|
||||
|
session_id = input_data.get("session_id", "unknown") |
||||
|
transcript_path = input_data.get("transcript_path", "") |
||||
|
cwd = input_data.get("cwd", "") |
||||
|
|
||||
|
log_file = get_log_file_path(session_id, cwd) |
||||
|
|
||||
|
if mode == "incremental": |
||||
|
# Stop event: incremental processing (fault tolerance) |
||||
|
messages = read_and_process_raw_transcript(transcript_path) |
||||
|
|
||||
|
if not messages: |
||||
|
return |
||||
|
|
||||
|
# Append all new unique messages to log |
||||
|
# Track what we've already logged (assistant by msg_id, thinking by msg_id, user by uuid) |
||||
|
existing_assistant_ids = set() |
||||
|
existing_thinking_ids = set() |
||||
|
existing_user_uuids = set() |
||||
|
|
||||
|
if os.path.exists(log_file): |
||||
|
with open(log_file, 'r', encoding='utf-8') as f: |
||||
|
for line in f: |
||||
|
try: |
||||
|
event = json.loads(line) |
||||
|
event_type = event.get('type') |
||||
|
if event_type == 'assistant': |
||||
|
msg = event.get('message', {}) |
||||
|
if msg.get('id'): |
||||
|
existing_assistant_ids.add(msg['id']) |
||||
|
elif event_type == 'assistant_thinking': |
||||
|
msg_id = event.get('message_id') |
||||
|
if msg_id: |
||||
|
existing_thinking_ids.add(msg_id) |
||||
|
elif event_type == 'user': |
||||
|
uuid = event.get('uuid') |
||||
|
if uuid: |
||||
|
existing_user_uuids.add(uuid) |
||||
|
except: |
||||
|
continue |
||||
|
|
||||
|
# Append new messages |
||||
|
new_count = 0 |
||||
|
with open(log_file, "a", encoding="utf-8") as f: |
||||
|
for msg_data in messages: |
||||
|
# Check if this is a new message |
||||
|
is_new = False |
||||
|
if msg_data['type'] == 'assistant': |
||||
|
msg_id = msg_data['message'].get('id') |
||||
|
if msg_id and msg_id not in existing_assistant_ids: |
||||
|
is_new = True |
||||
|
existing_assistant_ids.add(msg_id) |
||||
|
elif msg_data['type'] == 'assistant_thinking': |
||||
|
msg_id = msg_data.get('message_id') |
||||
|
if msg_id and msg_id not in existing_thinking_ids: |
||||
|
is_new = True |
||||
|
existing_thinking_ids.add(msg_id) |
||||
|
elif msg_data['type'] == 'user': |
||||
|
uuid = msg_data.get('uuid') |
||||
|
if uuid and uuid not in existing_user_uuids: |
||||
|
is_new = True |
||||
|
existing_user_uuids.add(uuid) |
||||
|
|
||||
|
if is_new: |
||||
|
# Add A/B metadata |
||||
|
log_entry = add_ab_metadata(msg_data.copy(), cwd) |
||||
|
f.write(json.dumps(log_entry) + "\n") |
||||
|
new_count += 1 |
||||
|
|
||||
|
if new_count > 0: |
||||
|
print(f"[OK] Processed {new_count} new messages (total: {len(messages)} unique)") |
||||
|
|
||||
|
elif mode == "final": |
||||
|
# SessionEnd: complete processing + summary |
||||
|
|
||||
|
# Step 1: Copy raw transcript |
||||
|
copy_raw_transcript(transcript_path, session_id, cwd) |
||||
|
|
||||
|
# Step 2: Process complete raw transcript |
||||
|
messages = read_and_process_raw_transcript(transcript_path) |
||||
|
|
||||
|
if not messages: |
||||
|
print("Warning: No messages found in raw transcript", file=sys.stderr) |
||||
|
return |
||||
|
|
||||
|
# Step 3: REBUILD processed log in perfect chronological order |
||||
|
# Read existing non-message events (session_start, etc.) |
||||
|
non_message_events = [] |
||||
|
|
||||
|
if os.path.exists(log_file): |
||||
|
with open(log_file, 'r', encoding='utf-8') as f: |
||||
|
for line in f: |
||||
|
try: |
||||
|
event = json.loads(line) |
||||
|
# Keep session_start and other non-message events |
||||
|
# Exclude assistant, assistant_thinking, and user messages (they come from raw transcript) |
||||
|
if event.get('type') not in ['assistant', 'assistant_thinking', 'user']: |
||||
|
non_message_events.append(event) |
||||
|
except: |
||||
|
continue |
||||
|
|
||||
|
# Combine all events and sort by timestamp |
||||
|
session_start = [e for e in non_message_events if e.get('type') == 'session_start'] |
||||
|
other_events = [e for e in non_message_events if e.get('type') != 'session_start'] |
||||
|
|
||||
|
# Build chronological list: session_start first, then messages sorted by time |
||||
|
all_events = [] |
||||
|
|
||||
|
# Add session_start first (if exists) |
||||
|
if session_start: |
||||
|
all_events.extend(session_start) |
||||
|
|
||||
|
# Add all messages (already sorted by timestamp from read_and_process_raw_transcript) |
||||
|
# Messages already include thinking blocks inserted before their parent assistant message |
||||
|
for msg_data in messages: |
||||
|
all_events.append(add_ab_metadata(msg_data.copy(), cwd)) |
||||
|
|
||||
|
print(f"[OK] Rebuilding log with {len(all_events)} events in chronological order") |
||||
|
|
||||
|
# Step 4: Generate session summary |
||||
|
usage_totals = aggregate_token_usage(messages) |
||||
|
tool_metrics = analyze_tool_calls(messages) |
||||
|
thinking_metrics = analyze_thinking_usage(messages, transcript_path) |
||||
|
|
||||
|
# Calculate duration |
||||
|
timestamps = [ |
||||
|
datetime.fromisoformat(msg['timestamp'].replace('Z', '+00:00')) |
||||
|
for msg in messages if msg.get('timestamp') |
||||
|
] |
||||
|
|
||||
|
total_duration = 0 |
||||
|
if len(timestamps) >= 2: |
||||
|
duration = max(timestamps) - min(timestamps) |
||||
|
total_duration = duration.total_seconds() |
||||
|
|
||||
|
# Count messages with proper categorization |
||||
|
# Note: assistant_thinking is NOT counted as a separate message (it's part of assistant message) |
||||
|
assistant_count = sum(1 for m in messages if m['type'] == 'assistant') |
||||
|
thinking_count = sum(1 for m in messages if m['type'] == 'assistant_thinking') |
||||
|
|
||||
|
# Categorize user messages |
||||
|
user_prompts = 0 |
||||
|
tool_results = 0 |
||||
|
system_messages = 0 |
||||
|
|
||||
|
for m in messages: |
||||
|
if m['type'] == 'user': |
||||
|
message_content = m['message'].get('content', '') |
||||
|
|
||||
|
# Check if it's a system/meta message |
||||
|
if m.get('isMeta'): |
||||
|
system_messages += 1 |
||||
|
# Check if it's a tool result |
||||
|
elif isinstance(message_content, list): |
||||
|
has_tool_result = any( |
||||
|
isinstance(item, dict) and item.get('type') == 'tool_result' |
||||
|
for item in message_content |
||||
|
) |
||||
|
if has_tool_result: |
||||
|
tool_results += 1 |
||||
|
else: |
||||
|
user_prompts += 1 |
||||
|
# Check if it's an exit/system command |
||||
|
elif isinstance(message_content, str) and ( |
||||
|
'<command-name>' in message_content or |
||||
|
'<local-command-stdout>' in message_content |
||||
|
): |
||||
|
system_messages += 1 |
||||
|
# Real user prompt (string content, not system) |
||||
|
elif isinstance(message_content, str): |
||||
|
user_prompts += 1 |
||||
|
else: |
||||
|
user_prompts += 1 # Default to user prompt |
||||
|
|
||||
|
total_user_events = user_prompts + tool_results + system_messages |
||||
|
|
||||
|
# Calculate actual total messages (excluding thinking blocks as they're not separate messages) |
||||
|
actual_total_messages = assistant_count + total_user_events |
||||
|
|
||||
|
# Get git metrics |
||||
|
base_commit = get_base_commit_from_log(log_file) |
||||
|
git_metrics = calculate_git_metrics(cwd, base_commit) if base_commit else {} |
||||
|
|
||||
|
# Create session summary |
||||
|
model_lane = detect_model_lane(cwd) |
||||
|
|
||||
|
summary = { |
||||
|
"type": "session_summary", |
||||
|
"timestamp": datetime.now(timezone.utc).isoformat(), |
||||
|
"session_id": session_id, |
||||
|
"transcript_path": transcript_path, |
||||
|
"cwd": cwd, |
||||
|
"summary_data": { |
||||
|
"total_duration_seconds": round(total_duration, 2), |
||||
|
"total_messages": actual_total_messages, |
||||
|
"assistant_messages": assistant_count, |
||||
|
"user_prompts": user_prompts, |
||||
|
"user_metrics": { |
||||
|
"user_prompts": user_prompts, |
||||
|
"tool_results": tool_results, |
||||
|
"system_messages": system_messages, |
||||
|
"total_user_events": total_user_events |
||||
|
}, |
||||
|
"usage_totals": usage_totals, |
||||
|
"tool_metrics": tool_metrics, |
||||
|
"thinking_metrics": { |
||||
|
**thinking_metrics, |
||||
|
"assistant_thinking_blocks_captured": thinking_count |
||||
|
}, |
||||
|
"git_metrics": git_metrics, |
||||
|
"files": { |
||||
|
"processed_log": f"session_{session_id}.jsonl", |
||||
|
"raw_transcript": f"session_{session_id}_raw.jsonl", |
||||
|
"git_diff": f"{model_lane}_diff.patch" if model_lane else None |
||||
|
}, |
||||
|
"validation": { |
||||
|
"complete": True, |
||||
|
"unique_messages_processed": actual_total_messages, |
||||
|
"thinking_blocks_extracted": thinking_count |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
summary = add_ab_metadata(summary, cwd) |
||||
|
|
||||
|
# Add session summary to events |
||||
|
all_events.append(summary) |
||||
|
|
||||
|
# Step 5: Rewrite log file with all events in perfect chronological order |
||||
|
# Write to temp file first, then rename (atomic) |
||||
|
temp_log_file = log_file + ".tmp" |
||||
|
|
||||
|
with open(temp_log_file, "w", encoding="utf-8") as f: |
||||
|
for event in all_events: |
||||
|
f.write(json.dumps(event) + "\n") |
||||
|
|
||||
|
# Atomic rename |
||||
|
os.replace(temp_log_file, log_file) |
||||
|
|
||||
|
print(f"[OK] Rebuilt log with {len(all_events)} events in chronological order") |
||||
|
print(f"[OK] Generated session summary: {actual_total_messages} messages, {assistant_count} assistant, {user_prompts} user prompts") |
||||
|
if thinking_count > 0: |
||||
|
print(f"[OK] Captured {thinking_count} thinking blocks (tokens already included in assistant output)") |
||||
|
print(f"[OK] User breakdown: {user_prompts} prompts, {tool_results} tool results, {system_messages} system") |
||||
|
print(f"[OK] Tokens: {usage_totals['total_actual_input_tokens']:,} input, {usage_totals['total_output_tokens']:,} output") |
||||
|
|
||||
|
except Exception as e: |
||||
|
print(f"[ERROR] Processing transcript: {e}", file=sys.stderr) |
||||
|
sys.exit(1) |
||||
|
|
||||
|
def calculate_git_metrics(cwd, base_commit): |
||||
|
"""Calculate git metrics from diff.""" |
||||
|
try: |
||||
|
original_cwd = os.getcwd() |
||||
|
os.chdir(cwd) |
||||
|
|
||||
|
if not base_commit: |
||||
|
os.chdir(original_cwd) |
||||
|
return {} |
||||
|
|
||||
|
# Add untracked files |
||||
|
excluded_patterns = ['.claude/', '__pycache__/', 'node_modules/', '.mypy_cache/', |
||||
|
'.pytest_cache/', '.DS_Store', '.vscode/', '.idea/'] |
||||
|
|
||||
|
untracked_result = subprocess.run( |
||||
|
['git', 'ls-files', '--others', '--exclude-standard'], |
||||
|
capture_output=True, text=True, timeout=30 |
||||
|
) |
||||
|
|
||||
|
if untracked_result.returncode == 0 and untracked_result.stdout.strip(): |
||||
|
untracked_files = [ |
||||
|
f.strip() for f in untracked_result.stdout.strip().split('\n') |
||||
|
if f.strip() and not any(pattern in f for pattern in excluded_patterns) |
||||
|
] |
||||
|
|
||||
|
for file in untracked_files: |
||||
|
subprocess.run(['git', 'add', '-N', file], capture_output=True, timeout=5) |
||||
|
|
||||
|
# Calculate numstat |
||||
|
result = subprocess.run( |
||||
|
['git', 'diff', '--numstat', base_commit, '--', '.', |
||||
|
':!.claude', ':!**/.mypy_cache', ':!**/__pycache__', ':!**/.pytest_cache', |
||||
|
':!**/.DS_Store', ':!**/node_modules', ':!**/.vscode', ':!**/.idea'], |
||||
|
capture_output=True, text=True, timeout=30 |
||||
|
) |
||||
|
|
||||
|
os.chdir(original_cwd) |
||||
|
|
||||
|
if result.returncode != 0: |
||||
|
return {} |
||||
|
|
||||
|
lines = result.stdout.strip().split('\n') if result.stdout.strip() else [] |
||||
|
files_changed = 0 |
||||
|
total_lines_changed = 0 |
||||
|
|
||||
|
for line in lines: |
||||
|
if line.strip(): |
||||
|
parts = line.split('\t') |
||||
|
if len(parts) >= 3: |
||||
|
try: |
||||
|
added = int(parts[0]) if parts[0] != '-' else 0 |
||||
|
removed = int(parts[1]) if parts[1] != '-' else 0 |
||||
|
files_changed += 1 |
||||
|
total_lines_changed += added + removed |
||||
|
except ValueError: |
||||
|
continue |
||||
|
|
||||
|
return { |
||||
|
"files_changed_count": files_changed, |
||||
|
"lines_of_code_changed_count": total_lines_changed |
||||
|
} |
||||
|
|
||||
|
except Exception as e: |
||||
|
print(f"Warning: Could not calculate git metrics: {e}", file=sys.stderr) |
||||
|
if 'original_cwd' in locals(): |
||||
|
os.chdir(original_cwd) |
||||
|
return {} |
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
main() |
||||
|
|
||||
@ -0,0 +1,55 @@ |
|||||
|
{ |
||||
|
"env": { |
||||
|
"ANTHROPIC_AUTH_TOKEN": "sk-1234-tap-tap-go-new", |
||||
|
"ANTHROPIC_BASE_URL": "https://mercor-rl--litellm-proxy-serve.modal.run" |
||||
|
}, |
||||
|
"permissions": { |
||||
|
"allow": [ |
||||
|
"Bash(ls:*)", |
||||
|
"Bash(grep:*)", |
||||
|
"Bash(python -m pytest:*)", |
||||
|
"Bash(pip install:*)", |
||||
|
"Bash(python3:*)", |
||||
|
"Bash(python:*)", |
||||
|
"WebSearch", |
||||
|
"WebFetch(domain:github.com)" |
||||
|
] |
||||
|
}, |
||||
|
"model": "zorilla", |
||||
|
"hooks": { |
||||
|
"SessionStart": [ |
||||
|
{ |
||||
|
"hooks": [ |
||||
|
{ |
||||
|
"type": "command", |
||||
|
"command": "python \"$CLAUDE_PROJECT_DIR/.claude/hooks/capture_session_event.py\" start" |
||||
|
} |
||||
|
] |
||||
|
} |
||||
|
], |
||||
|
"Stop": [ |
||||
|
{ |
||||
|
"hooks": [ |
||||
|
{ |
||||
|
"type": "command", |
||||
|
"command": "python \"$CLAUDE_PROJECT_DIR/.claude/hooks/process_transcript.py\" incremental" |
||||
|
} |
||||
|
] |
||||
|
} |
||||
|
], |
||||
|
"SessionEnd": [ |
||||
|
{ |
||||
|
"hooks": [ |
||||
|
{ |
||||
|
"type": "command", |
||||
|
"command": "python \"$CLAUDE_PROJECT_DIR/.claude/hooks/process_transcript.py\" final" |
||||
|
}, |
||||
|
{ |
||||
|
"type": "command", |
||||
|
"command": "python \"$CLAUDE_PROJECT_DIR/.claude/hooks/capture_session_event.py\" end" |
||||
|
} |
||||
|
] |
||||
|
} |
||||
|
] |
||||
|
} |
||||
|
} |
||||
@ -0,0 +1,203 @@ |
|||||
|
"""Unified lifespan management for FastAPI.""" |
||||
|
|
||||
|
import warnings |
||||
|
from collections.abc import AsyncGenerator, Callable, Sequence |
||||
|
from contextlib import asynccontextmanager |
||||
|
from typing import Any, Optional, TypeVar |
||||
|
|
||||
|
from starlette.types import ASGIApp, Lifespan |
||||
|
|
||||
|
from fastapi.exceptions import FastAPIDeprecationWarning |
||||
|
|
||||
|
AppType = TypeVar("AppType", bound="ASGIApp") |
||||
|
|
||||
|
|
||||
|
class LifespanError(Exception): |
||||
|
"""Exception raised when lifespan handling fails.""" |
||||
|
|
||||
|
def __init__( |
||||
|
self, |
||||
|
message: str, |
||||
|
*, |
||||
|
phase: str, |
||||
|
handler: Optional[Callable[..., Any]] = None, |
||||
|
original_error: Optional[BaseException] = None, |
||||
|
): |
||||
|
super().__init__(message) |
||||
|
self.phase = phase |
||||
|
self.handler = handler |
||||
|
self.original_error = original_error |
||||
|
|
||||
|
def __str__(self) -> str: |
||||
|
base = f"Lifespan {self.phase} failed" |
||||
|
if self.handler: |
||||
|
handler_name = getattr(self.handler, "__name__", str(self.handler)) |
||||
|
base += f" in handler '{handler_name}'" |
||||
|
if self.original_error: |
||||
|
base += f": {self.original_error}" |
||||
|
return base |
||||
|
|
||||
|
|
||||
|
def _emit_on_event_deprecation_warning(event_type: str) -> None: |
||||
|
warnings.warn( |
||||
|
f"on_event('{event_type}') is deprecated. Use the lifespan context manager.", |
||||
|
FastAPIDeprecationWarning, |
||||
|
stacklevel=4, |
||||
|
) |
||||
|
|
||||
|
|
||||
|
class UnifiedLifespanManager: |
||||
|
"""Manages unified lifespan handling with rollback semantics.""" |
||||
|
|
||||
|
def __init__( |
||||
|
self, |
||||
|
*, |
||||
|
on_startup: Optional[Sequence[Callable[..., Any]]] = None, |
||||
|
on_shutdown: Optional[Sequence[Callable[..., Any]]] = None, |
||||
|
lifespan: Optional[Lifespan[AppType]] = None, |
||||
|
emit_deprecation_warnings: bool = True, |
||||
|
): |
||||
|
self._on_startup = list(on_startup or []) |
||||
|
self._on_shutdown = list(on_shutdown or []) |
||||
|
self._lifespan = lifespan |
||||
|
self._emit_warnings = emit_deprecation_warnings |
||||
|
self._completed_startup_handlers: list[Callable[..., Any]] = [] |
||||
|
|
||||
|
if lifespan and (on_startup or on_shutdown): |
||||
|
warnings.warn( |
||||
|
"Using both 'lifespan' and 'on_startup'/'on_shutdown' is not recommended.", |
||||
|
FastAPIDeprecationWarning, |
||||
|
stacklevel=2, |
||||
|
) |
||||
|
|
||||
|
@asynccontextmanager |
||||
|
async def lifespan(self, app: AppType) -> AsyncGenerator[dict[str, Any], None]: |
||||
|
"""Unified lifespan context manager with rollback semantics.""" |
||||
|
state: dict[str, Any] = {} |
||||
|
self._completed_startup_handlers = [] |
||||
|
lifespan_cm = None |
||||
|
|
||||
|
try: |
||||
|
await self._execute_startup_handlers() |
||||
|
|
||||
|
if self._lifespan: |
||||
|
lifespan_cm = self._lifespan(app) |
||||
|
lifespan_state = await lifespan_cm.__aenter__() |
||||
|
if lifespan_state: |
||||
|
state.update(lifespan_state) |
||||
|
|
||||
|
try: |
||||
|
yield state |
||||
|
finally: |
||||
|
shutdown_errors: list[tuple[str, BaseException]] = [] |
||||
|
|
||||
|
if lifespan_cm is not None: |
||||
|
try: |
||||
|
await lifespan_cm.__aexit__(None, None, None) |
||||
|
except Exception as e: |
||||
|
shutdown_errors.append(("lifespan context", e)) |
||||
|
|
||||
|
handler_errors = await self._execute_shutdown_handlers() |
||||
|
shutdown_errors.extend(handler_errors) |
||||
|
|
||||
|
if shutdown_errors: |
||||
|
error_msgs = [f"{name}: {err}" for name, err in shutdown_errors] |
||||
|
warnings.warn( |
||||
|
f"Errors during shutdown: {'; '.join(error_msgs)}", |
||||
|
RuntimeWarning, |
||||
|
stacklevel=2, |
||||
|
) |
||||
|
|
||||
|
except BaseException as startup_error: |
||||
|
await self._rollback_startup(startup_error) |
||||
|
raise |
||||
|
|
||||
|
async def _execute_startup_handlers(self) -> None: |
||||
|
for handler in self._on_startup: |
||||
|
if self._emit_warnings: |
||||
|
_emit_on_event_deprecation_warning("startup") |
||||
|
|
||||
|
try: |
||||
|
result = handler() |
||||
|
if hasattr(result, "__await__"): |
||||
|
await result |
||||
|
self._completed_startup_handlers.append(handler) |
||||
|
except Exception as e: |
||||
|
raise LifespanError( |
||||
|
f"Startup handler failed", |
||||
|
phase="startup", |
||||
|
handler=handler, |
||||
|
original_error=e, |
||||
|
) from e |
||||
|
|
||||
|
async def _execute_shutdown_handlers(self) -> list[tuple[str, BaseException]]: |
||||
|
errors: list[tuple[str, BaseException]] = [] |
||||
|
for handler in reversed(self._on_shutdown): |
||||
|
if self._emit_warnings: |
||||
|
_emit_on_event_deprecation_warning("shutdown") |
||||
|
|
||||
|
try: |
||||
|
result = handler() |
||||
|
if hasattr(result, "__await__"): |
||||
|
await result |
||||
|
except Exception as e: |
||||
|
handler_name = getattr(handler, "__name__", str(handler)) |
||||
|
errors.append((handler_name, e)) |
||||
|
return errors |
||||
|
|
||||
|
async def _rollback_startup(self, original_error: BaseException) -> None: |
||||
|
if not self._completed_startup_handlers: |
||||
|
return |
||||
|
|
||||
|
completed_count = len(self._completed_startup_handlers) |
||||
|
handlers_to_cleanup = list(reversed(self._on_shutdown[:completed_count])) |
||||
|
|
||||
|
for handler in handlers_to_cleanup: |
||||
|
try: |
||||
|
result = handler() |
||||
|
if hasattr(result, "__await__"): |
||||
|
await result |
||||
|
except Exception as cleanup_error: |
||||
|
handler_name = getattr(handler, "__name__", str(handler)) |
||||
|
warnings.warn( |
||||
|
f"Error during rollback in '{handler_name}': {cleanup_error}", |
||||
|
RuntimeWarning, |
||||
|
stacklevel=2, |
||||
|
) |
||||
|
|
||||
|
def add_startup_handler(self, handler: Callable[..., Any]) -> None: |
||||
|
self._on_startup.append(handler) |
||||
|
|
||||
|
def add_shutdown_handler(self, handler: Callable[..., Any]) -> None: |
||||
|
self._on_shutdown.append(handler) |
||||
|
|
||||
|
|
||||
|
def create_unified_lifespan( |
||||
|
*, |
||||
|
on_startup: Optional[Sequence[Callable[..., Any]]] = None, |
||||
|
on_shutdown: Optional[Sequence[Callable[..., Any]]] = None, |
||||
|
lifespan: Optional[Lifespan[AppType]] = None, |
||||
|
) -> Lifespan[AppType]: |
||||
|
"""Create a unified lifespan from mixed inputs.""" |
||||
|
manager = UnifiedLifespanManager( |
||||
|
on_startup=on_startup, |
||||
|
on_shutdown=on_shutdown, |
||||
|
lifespan=lifespan, |
||||
|
) |
||||
|
return manager.lifespan # type: ignore |
||||
|
|
||||
|
|
||||
|
def merge_lifespan_state( |
||||
|
parent_state: Optional[dict[str, Any]], |
||||
|
child_state: Optional[dict[str, Any]], |
||||
|
*, |
||||
|
child_overrides_parent: bool = True, |
||||
|
) -> dict[str, Any]: |
||||
|
"""Merge lifespan state dictionaries.""" |
||||
|
parent = parent_state or {} |
||||
|
child = child_state or {} |
||||
|
|
||||
|
if child_overrides_parent: |
||||
|
return {**parent, **child} |
||||
|
else: |
||||
|
return {**child, **parent} |
||||
@ -0,0 +1,205 @@ |
|||||
|
"""OpenAPI schema generation profiling utilities.""" |
||||
|
|
||||
|
import functools |
||||
|
import threading |
||||
|
import time |
||||
|
from collections import defaultdict |
||||
|
from contextlib import contextmanager |
||||
|
from dataclasses import dataclass, field |
||||
|
from typing import Any, Callable, Optional, TypeVar |
||||
|
|
||||
|
F = TypeVar("F", bound=Callable[..., Any]) |
||||
|
|
||||
|
|
||||
|
@dataclass |
||||
|
class TimingEntry: |
||||
|
"""Single timing measurement entry.""" |
||||
|
|
||||
|
name: str |
||||
|
duration_ms: float |
||||
|
call_count: int = 1 |
||||
|
metadata: dict[str, Any] = field(default_factory=dict) |
||||
|
|
||||
|
|
||||
|
@dataclass |
||||
|
class ProfilingStats: |
||||
|
"""Aggregated profiling statistics for a function.""" |
||||
|
|
||||
|
name: str |
||||
|
total_time_ms: float = 0.0 |
||||
|
call_count: int = 0 |
||||
|
min_time_ms: float = float("inf") |
||||
|
max_time_ms: float = 0.0 |
||||
|
timings: list[float] = field(default_factory=list) |
||||
|
|
||||
|
@property |
||||
|
def avg_time_ms(self) -> float: |
||||
|
return self.total_time_ms / self.call_count if self.call_count > 0 else 0.0 |
||||
|
|
||||
|
def add_timing(self, duration_ms: float) -> None: |
||||
|
self.total_time_ms += duration_ms |
||||
|
self.call_count += 1 |
||||
|
self.min_time_ms = min(self.min_time_ms, duration_ms) |
||||
|
self.max_time_ms = max(self.max_time_ms, duration_ms) |
||||
|
self.timings.append(duration_ms) |
||||
|
|
||||
|
|
||||
|
class OpenAPIProfiler: |
||||
|
"""Thread-safe profiler for OpenAPI schema generation.""" |
||||
|
|
||||
|
def __init__(self) -> None: |
||||
|
self._enabled = False |
||||
|
self._stats: dict[str, ProfilingStats] = defaultdict( |
||||
|
lambda: ProfilingStats(name="") |
||||
|
) |
||||
|
self._lock = threading.Lock() |
||||
|
self._context_stack: list[str] = [] |
||||
|
self._start_time: Optional[float] = None |
||||
|
|
||||
|
def enable(self) -> None: |
||||
|
"""Enable profiling.""" |
||||
|
with self._lock: |
||||
|
self._enabled = True |
||||
|
self._start_time = time.perf_counter() |
||||
|
|
||||
|
def disable(self) -> None: |
||||
|
"""Disable profiling.""" |
||||
|
with self._lock: |
||||
|
self._enabled = False |
||||
|
|
||||
|
def reset(self) -> None: |
||||
|
"""Reset all collected statistics.""" |
||||
|
with self._lock: |
||||
|
self._stats.clear() |
||||
|
self._context_stack.clear() |
||||
|
self._start_time = None |
||||
|
|
||||
|
@property |
||||
|
def is_enabled(self) -> bool: |
||||
|
return self._enabled |
||||
|
|
||||
|
def record(self, name: str, duration_ms: float) -> None: |
||||
|
"""Record a timing measurement.""" |
||||
|
if not self._enabled: |
||||
|
return |
||||
|
with self._lock: |
||||
|
if name not in self._stats: |
||||
|
self._stats[name] = ProfilingStats(name=name) |
||||
|
self._stats[name].add_timing(duration_ms) |
||||
|
|
||||
|
@contextmanager |
||||
|
def measure(self, name: str): |
||||
|
"""Context manager to measure execution time of a code block.""" |
||||
|
if not self._enabled: |
||||
|
yield |
||||
|
return |
||||
|
|
||||
|
start = time.perf_counter() |
||||
|
try: |
||||
|
yield |
||||
|
finally: |
||||
|
duration_ms = (time.perf_counter() - start) * 1000 |
||||
|
self.record(name, duration_ms) |
||||
|
|
||||
|
def get_stats(self) -> dict[str, ProfilingStats]: |
||||
|
"""Get a copy of current statistics.""" |
||||
|
with self._lock: |
||||
|
return dict(self._stats) |
||||
|
|
||||
|
def get_total_time_ms(self) -> float: |
||||
|
"""Get total elapsed time since profiling was enabled.""" |
||||
|
if self._start_time is None: |
||||
|
return 0.0 |
||||
|
return (time.perf_counter() - self._start_time) * 1000 |
||||
|
|
||||
|
def get_report(self) -> str: |
||||
|
"""Generate a human-readable profiling report.""" |
||||
|
stats = self.get_stats() |
||||
|
if not stats: |
||||
|
return "No profiling data collected." |
||||
|
|
||||
|
total_time = self.get_total_time_ms() |
||||
|
lines = [ |
||||
|
"=" * 80, |
||||
|
"OpenAPI Schema Generation Profiling Report", |
||||
|
"=" * 80, |
||||
|
f"Total elapsed time: {total_time:.2f}ms", |
||||
|
"", |
||||
|
f"{'Function':<45} {'Calls':>6} {'Total':>10} {'Avg':>10} {'%':>6}", |
||||
|
"-" * 80, |
||||
|
] |
||||
|
|
||||
|
sorted_stats = sorted( |
||||
|
stats.values(), key=lambda s: s.total_time_ms, reverse=True |
||||
|
) |
||||
|
|
||||
|
for stat in sorted_stats: |
||||
|
pct = (stat.total_time_ms / total_time * 100) if total_time > 0 else 0 |
||||
|
lines.append( |
||||
|
f"{stat.name:<45} {stat.call_count:>6} " |
||||
|
f"{stat.total_time_ms:>9.2f}ms {stat.avg_time_ms:>9.2f}ms {pct:>5.1f}%" |
||||
|
) |
||||
|
|
||||
|
lines.append("=" * 80) |
||||
|
return "\n".join(lines) |
||||
|
|
||||
|
def print_report(self) -> None: |
||||
|
"""Print the profiling report to stdout.""" |
||||
|
print(self.get_report()) |
||||
|
|
||||
|
|
||||
|
# Global profiler instance |
||||
|
openapi_profiler = OpenAPIProfiler() |
||||
|
|
||||
|
|
||||
|
def profiled(name: Optional[str] = None) -> Callable[[F], F]: |
||||
|
"""Decorator to instrument a function for profiling.""" |
||||
|
|
||||
|
def decorator(func: F) -> F: |
||||
|
func_name = name or f"{func.__module__}.{func.__qualname__}" |
||||
|
|
||||
|
@functools.wraps(func) |
||||
|
def wrapper(*args: Any, **kwargs: Any) -> Any: |
||||
|
if not openapi_profiler.is_enabled: |
||||
|
return func(*args, **kwargs) |
||||
|
|
||||
|
start = time.perf_counter() |
||||
|
try: |
||||
|
return func(*args, **kwargs) |
||||
|
finally: |
||||
|
duration_ms = (time.perf_counter() - start) * 1000 |
||||
|
openapi_profiler.record(func_name, duration_ms) |
||||
|
|
||||
|
return wrapper # type: ignore[return-value] |
||||
|
|
||||
|
return decorator |
||||
|
|
||||
|
|
||||
|
class ProfilingContext: |
||||
|
"""Context manager for scoped profiling sessions.""" |
||||
|
|
||||
|
def __init__(self, auto_print: bool = False) -> None: |
||||
|
self._auto_print = auto_print |
||||
|
self._profiler = openapi_profiler |
||||
|
|
||||
|
def __enter__(self) -> "ProfilingContext": |
||||
|
self._profiler.reset() |
||||
|
self._profiler.enable() |
||||
|
return self |
||||
|
|
||||
|
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: |
||||
|
self._profiler.disable() |
||||
|
if self._auto_print: |
||||
|
self.print_report() |
||||
|
|
||||
|
def get_stats(self) -> dict[str, ProfilingStats]: |
||||
|
return self._profiler.get_stats() |
||||
|
|
||||
|
def get_report(self) -> str: |
||||
|
return self._profiler.get_report() |
||||
|
|
||||
|
def print_report(self) -> None: |
||||
|
self._profiler.print_report() |
||||
|
|
||||
|
def get_total_time_ms(self) -> float: |
||||
|
return self._profiler.get_total_time_ms() |
||||
@ -0,0 +1,376 @@ |
|||||
|
"""OpenAPI plugin system for extending schema generation.""" |
||||
|
|
||||
|
import threading |
||||
|
import warnings |
||||
|
from abc import ABC, abstractmethod |
||||
|
from dataclasses import dataclass |
||||
|
from typing import TYPE_CHECKING, Any, Optional, Protocol, Union, runtime_checkable |
||||
|
|
||||
|
if TYPE_CHECKING: |
||||
|
from fastapi import FastAPI |
||||
|
from fastapi.routing import APIRoute |
||||
|
|
||||
|
|
||||
|
@dataclass |
||||
|
class OpenAPIPluginSettings: |
||||
|
"""Settings passed to plugins during schema generation.""" |
||||
|
|
||||
|
title: str |
||||
|
version: str |
||||
|
openapi_version: str |
||||
|
description: Optional[str] = None |
||||
|
routes_count: int = 0 |
||||
|
separate_input_output_schemas: bool = True |
||||
|
|
||||
|
|
||||
|
@runtime_checkable |
||||
|
class OpenAPIPlugin(Protocol): |
||||
|
"""Protocol defining the interface for OpenAPI plugins.""" |
||||
|
|
||||
|
@property |
||||
|
def name(self) -> str: |
||||
|
"""Unique identifier for the plugin.""" |
||||
|
... |
||||
|
|
||||
|
@property |
||||
|
def priority(self) -> int: |
||||
|
"""Execution priority. Lower numbers execute first. Default is 100.""" |
||||
|
... |
||||
|
|
||||
|
def pre_schema_generation( |
||||
|
self, |
||||
|
app: "FastAPI", |
||||
|
settings: OpenAPIPluginSettings, |
||||
|
) -> None: |
||||
|
"""Called before schema generation begins.""" |
||||
|
... |
||||
|
|
||||
|
def modify_operation( |
||||
|
self, |
||||
|
route: "APIRoute", |
||||
|
method: str, |
||||
|
operation: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
"""Called for each operation. Returns the modified operation dict.""" |
||||
|
... |
||||
|
|
||||
|
def modify_schema( |
||||
|
self, |
||||
|
schema: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
"""Called after base schema is built. Returns the modified schema.""" |
||||
|
... |
||||
|
|
||||
|
def post_schema_generation( |
||||
|
self, |
||||
|
schema: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
"""Called as the final step. Returns the final schema.""" |
||||
|
... |
||||
|
|
||||
|
|
||||
|
class OpenAPIPluginBase(ABC): |
||||
|
"""Abstract base class for OpenAPI plugins with default implementations.""" |
||||
|
|
||||
|
@property |
||||
|
@abstractmethod |
||||
|
def name(self) -> str: |
||||
|
"""Unique plugin identifier.""" |
||||
|
raise NotImplementedError |
||||
|
|
||||
|
@property |
||||
|
def priority(self) -> int: |
||||
|
"""Execution priority. Override to change order.""" |
||||
|
return 100 |
||||
|
|
||||
|
def pre_schema_generation( |
||||
|
self, |
||||
|
app: "FastAPI", |
||||
|
settings: OpenAPIPluginSettings, |
||||
|
) -> None: |
||||
|
pass |
||||
|
|
||||
|
def modify_operation( |
||||
|
self, |
||||
|
route: "APIRoute", |
||||
|
method: str, |
||||
|
operation: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
return operation |
||||
|
|
||||
|
def modify_schema( |
||||
|
self, |
||||
|
schema: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
return schema |
||||
|
|
||||
|
def post_schema_generation( |
||||
|
self, |
||||
|
schema: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
return schema |
||||
|
|
||||
|
|
||||
|
@dataclass |
||||
|
class PluginState: |
||||
|
"""Internal state for a registered plugin.""" |
||||
|
|
||||
|
plugin: Union[OpenAPIPlugin, OpenAPIPluginBase] |
||||
|
enabled: bool = True |
||||
|
error_count: int = 0 |
||||
|
last_error: Optional[Exception] = None |
||||
|
|
||||
|
|
||||
|
class PluginRegistry: |
||||
|
"""Thread-safe registry for managing OpenAPI plugins.""" |
||||
|
|
||||
|
def __init__( |
||||
|
self, |
||||
|
*, |
||||
|
max_errors_before_disable: int = 3, |
||||
|
on_cache_invalidation: Optional[callable] = None, |
||||
|
) -> None: |
||||
|
self._plugins: dict[str, PluginState] = {} |
||||
|
self._lock = threading.RLock() |
||||
|
self._max_errors = max_errors_before_disable |
||||
|
self._cache_invalidation_callback = on_cache_invalidation |
||||
|
self._generation_count = 0 |
||||
|
|
||||
|
def register( |
||||
|
self, |
||||
|
plugin: Union[OpenAPIPlugin, OpenAPIPluginBase], |
||||
|
*, |
||||
|
enabled: bool = True, |
||||
|
replace: bool = False, |
||||
|
) -> None: |
||||
|
"""Register a plugin. Raises ValueError if name exists and replace=False.""" |
||||
|
if not hasattr(plugin, "name"): |
||||
|
raise TypeError( |
||||
|
f"Plugin must have a 'name' property, got {type(plugin).__name__}" |
||||
|
) |
||||
|
|
||||
|
name = plugin.name |
||||
|
if not isinstance(name, str) or not name: |
||||
|
raise ValueError(f"Plugin name must be a non-empty string, got {name!r}") |
||||
|
|
||||
|
with self._lock: |
||||
|
if name in self._plugins and not replace: |
||||
|
raise ValueError( |
||||
|
f"Plugin '{name}' is already registered. " |
||||
|
f"Use replace=True to override or unregister first." |
||||
|
) |
||||
|
self._plugins[name] = PluginState(plugin=plugin, enabled=enabled) |
||||
|
self._invalidate_cache() |
||||
|
|
||||
|
def unregister(self, name: str) -> bool: |
||||
|
"""Remove a plugin. Returns True if removed, False if not found.""" |
||||
|
with self._lock: |
||||
|
if name in self._plugins: |
||||
|
del self._plugins[name] |
||||
|
self._invalidate_cache() |
||||
|
return True |
||||
|
return False |
||||
|
|
||||
|
def enable(self, name: str) -> bool: |
||||
|
"""Enable a plugin. Returns True if enabled, False if not found.""" |
||||
|
with self._lock: |
||||
|
if name in self._plugins: |
||||
|
state = self._plugins[name] |
||||
|
if not state.enabled: |
||||
|
state.enabled = True |
||||
|
state.error_count = 0 |
||||
|
self._invalidate_cache() |
||||
|
return True |
||||
|
return False |
||||
|
|
||||
|
def disable(self, name: str) -> bool: |
||||
|
"""Disable a plugin. Returns True if disabled, False if not found.""" |
||||
|
with self._lock: |
||||
|
if name in self._plugins: |
||||
|
state = self._plugins[name] |
||||
|
if state.enabled: |
||||
|
state.enabled = False |
||||
|
self._invalidate_cache() |
||||
|
return True |
||||
|
return False |
||||
|
|
||||
|
def is_enabled(self, name: str) -> bool: |
||||
|
"""Check if a plugin is enabled.""" |
||||
|
with self._lock: |
||||
|
state = self._plugins.get(name) |
||||
|
return state.enabled if state else False |
||||
|
|
||||
|
def get_plugin( |
||||
|
self, name: str |
||||
|
) -> Optional[Union[OpenAPIPlugin, OpenAPIPluginBase]]: |
||||
|
"""Get a plugin by name.""" |
||||
|
with self._lock: |
||||
|
state = self._plugins.get(name) |
||||
|
return state.plugin if state else None |
||||
|
|
||||
|
def get_active_plugins(self) -> list[Union[OpenAPIPlugin, OpenAPIPluginBase]]: |
||||
|
"""Get all enabled plugins sorted by priority (lowest first).""" |
||||
|
with self._lock: |
||||
|
active = [ |
||||
|
state.plugin for state in self._plugins.values() if state.enabled |
||||
|
] |
||||
|
return sorted(active, key=lambda p: getattr(p, "priority", 100)) |
||||
|
|
||||
|
def get_all_plugins( |
||||
|
self, |
||||
|
) -> dict[str, tuple[Union[OpenAPIPlugin, OpenAPIPluginBase], bool]]: |
||||
|
"""Get all plugins with their enabled status.""" |
||||
|
with self._lock: |
||||
|
return { |
||||
|
name: (state.plugin, state.enabled) |
||||
|
for name, state in self._plugins.items() |
||||
|
} |
||||
|
|
||||
|
def record_error(self, name: str, error: Exception) -> None: |
||||
|
"""Record an error. Auto-disables plugin after threshold.""" |
||||
|
with self._lock: |
||||
|
state = self._plugins.get(name) |
||||
|
if state: |
||||
|
state.error_count += 1 |
||||
|
state.last_error = error |
||||
|
if state.error_count >= self._max_errors: |
||||
|
warnings.warn( |
||||
|
f"OpenAPI plugin '{name}' auto-disabled after " |
||||
|
f"{state.error_count} errors. Last error: {error}", |
||||
|
stacklevel=2, |
||||
|
) |
||||
|
state.enabled = False |
||||
|
self._invalidate_cache() |
||||
|
|
||||
|
def clear(self) -> None: |
||||
|
"""Remove all plugins.""" |
||||
|
with self._lock: |
||||
|
self._plugins.clear() |
||||
|
self._invalidate_cache() |
||||
|
|
||||
|
@property |
||||
|
def generation_count(self) -> int: |
||||
|
"""Counter incremented on cache invalidation.""" |
||||
|
return self._generation_count |
||||
|
|
||||
|
def _invalidate_cache(self) -> None: |
||||
|
self._generation_count += 1 |
||||
|
if self._cache_invalidation_callback: |
||||
|
try: |
||||
|
self._cache_invalidation_callback() |
||||
|
except Exception: |
||||
|
pass |
||||
|
|
||||
|
def __len__(self) -> int: |
||||
|
return len(self._plugins) |
||||
|
|
||||
|
def __contains__(self, name: str) -> bool: |
||||
|
return name in self._plugins |
||||
|
|
||||
|
|
||||
|
class PluginExecutor: |
||||
|
"""Executes plugin hooks during schema generation.""" |
||||
|
|
||||
|
def __init__(self, registry: PluginRegistry) -> None: |
||||
|
self._registry = registry |
||||
|
|
||||
|
def execute_pre_schema_generation( |
||||
|
self, |
||||
|
app: "FastAPI", |
||||
|
settings: OpenAPIPluginSettings, |
||||
|
) -> None: |
||||
|
for plugin in self._registry.get_active_plugins(): |
||||
|
try: |
||||
|
if hasattr(plugin, "pre_schema_generation"): |
||||
|
plugin.pre_schema_generation(app, settings) |
||||
|
except Exception as e: |
||||
|
self._handle_error(plugin, "pre_schema_generation", e) |
||||
|
|
||||
|
def execute_modify_operation( |
||||
|
self, |
||||
|
route: "APIRoute", |
||||
|
method: str, |
||||
|
operation: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
result = operation |
||||
|
for plugin in self._registry.get_active_plugins(): |
||||
|
try: |
||||
|
if hasattr(plugin, "modify_operation"): |
||||
|
result = plugin.modify_operation(route, method, result) |
||||
|
if result is None: |
||||
|
result = operation |
||||
|
except Exception as e: |
||||
|
self._handle_error(plugin, "modify_operation", e) |
||||
|
return result |
||||
|
|
||||
|
def execute_modify_schema( |
||||
|
self, |
||||
|
schema: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
result = schema |
||||
|
for plugin in self._registry.get_active_plugins(): |
||||
|
try: |
||||
|
if hasattr(plugin, "modify_schema"): |
||||
|
result = plugin.modify_schema(result) |
||||
|
if result is None: |
||||
|
result = schema |
||||
|
except Exception as e: |
||||
|
self._handle_error(plugin, "modify_schema", e) |
||||
|
return result |
||||
|
|
||||
|
def execute_post_schema_generation( |
||||
|
self, |
||||
|
schema: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
result = schema |
||||
|
for plugin in self._registry.get_active_plugins(): |
||||
|
try: |
||||
|
if hasattr(plugin, "post_schema_generation"): |
||||
|
result = plugin.post_schema_generation(result) |
||||
|
if result is None: |
||||
|
result = schema |
||||
|
except Exception as e: |
||||
|
self._handle_error(plugin, "post_schema_generation", e) |
||||
|
return result |
||||
|
|
||||
|
def _handle_error(self, plugin: Any, hook_name: str, error: Exception) -> None: |
||||
|
plugin_name = getattr(plugin, "name", type(plugin).__name__) |
||||
|
warnings.warn( |
||||
|
f"OpenAPI plugin '{plugin_name}' raised an error in {hook_name}: {error}", |
||||
|
stacklevel=3, |
||||
|
) |
||||
|
self._registry.record_error(plugin_name, error) |
||||
|
|
||||
|
|
||||
|
def create_plugin( |
||||
|
name: str, |
||||
|
*, |
||||
|
priority: int = 100, |
||||
|
pre_schema_generation: Optional[callable] = None, |
||||
|
modify_operation: Optional[callable] = None, |
||||
|
modify_schema: Optional[callable] = None, |
||||
|
post_schema_generation: Optional[callable] = None, |
||||
|
) -> OpenAPIPluginBase: |
||||
|
"""Create a plugin from callback functions.""" |
||||
|
|
||||
|
class CallbackPlugin(OpenAPIPluginBase): |
||||
|
@property |
||||
|
def name(self) -> str: |
||||
|
return name |
||||
|
|
||||
|
@property |
||||
|
def priority(self) -> int: |
||||
|
return priority |
||||
|
|
||||
|
plugin = CallbackPlugin() |
||||
|
|
||||
|
if pre_schema_generation: |
||||
|
plugin.pre_schema_generation = pre_schema_generation # type: ignore |
||||
|
if modify_operation: |
||||
|
plugin.modify_operation = modify_operation # type: ignore |
||||
|
if modify_schema: |
||||
|
plugin.modify_schema = modify_schema # type: ignore |
||||
|
if post_schema_generation: |
||||
|
plugin.post_schema_generation = post_schema_generation # type: ignore |
||||
|
|
||||
|
return plugin |
||||
File diff suppressed because it is too large
@ -0,0 +1,370 @@ |
|||||
|
""" |
||||
|
Test for potential race conditions in async dependencies with BackgroundTasks. |
||||
|
|
||||
|
These tests specifically look for: |
||||
|
1. Race conditions under concurrent requests |
||||
|
2. Cleanup order issues with nested dependencies |
||||
|
3. Exception handling during cleanup |
||||
|
4. Resource access after cleanup starts |
||||
|
""" |
||||
|
|
||||
|
import asyncio |
||||
|
import threading |
||||
|
from collections import defaultdict |
||||
|
from typing import Any |
||||
|
|
||||
|
import pytest |
||||
|
from fastapi import BackgroundTasks, Depends, FastAPI, HTTPException |
||||
|
from fastapi.testclient import TestClient |
||||
|
|
||||
|
|
||||
|
# Test 1: Verify cleanup timing under simulated high concurrency |
||||
|
def test_cleanup_timing_consistency(): |
||||
|
""" |
||||
|
Test that cleanup timing is consistent across multiple requests. |
||||
|
The cleanup should always happen at the same point relative to background tasks. |
||||
|
""" |
||||
|
results: list[dict] = [] |
||||
|
lock = threading.Lock() |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_resource(): |
||||
|
resource = {"state": "active", "cleanup_happened": False} |
||||
|
yield resource |
||||
|
resource["cleanup_happened"] = True |
||||
|
resource["state"] = "cleaned" |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: dict = Depends(get_resource), |
||||
|
): |
||||
|
request_id = len(results) |
||||
|
|
||||
|
async def bg_task(res: dict, req_id: int): |
||||
|
# Record state when background task runs |
||||
|
with lock: |
||||
|
results.append({ |
||||
|
"request_id": req_id, |
||||
|
"state_in_bg": res["state"], |
||||
|
"cleanup_in_bg": res["cleanup_happened"], |
||||
|
}) |
||||
|
|
||||
|
background_tasks.add_task(bg_task, resource, request_id) |
||||
|
return {"request_id": request_id} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
|
||||
|
# Run multiple requests |
||||
|
for _ in range(20): |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# All results should be consistent - with default scope, cleanup should |
||||
|
# NOT have happened when background task runs |
||||
|
for result in results: |
||||
|
assert result["state_in_bg"] == "active", f"Inconsistent state: {result}" |
||||
|
assert result["cleanup_in_bg"] is False, f"Cleanup happened too early: {result}" |
||||
|
|
||||
|
|
||||
|
# Test 2: Test for cleanup happening during background task execution |
||||
|
def test_cleanup_during_background_task_execution(): |
||||
|
""" |
||||
|
Test if cleanup can start while a background task is still running. |
||||
|
This shouldn't happen with proper request-scoped dependencies. |
||||
|
""" |
||||
|
events: list[str] = [] |
||||
|
lock = threading.Lock() |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_slow_resource(): |
||||
|
with lock: |
||||
|
events.append("resource_start") |
||||
|
yield "resource" |
||||
|
with lock: |
||||
|
events.append("resource_cleanup") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: str = Depends(get_slow_resource), |
||||
|
): |
||||
|
async def slow_bg_task(): |
||||
|
with lock: |
||||
|
events.append("bg_start") |
||||
|
await asyncio.sleep(0.05) # Simulate slow work |
||||
|
with lock: |
||||
|
events.append("bg_end") |
||||
|
|
||||
|
background_tasks.add_task(slow_bg_task) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Verify order: resource_start -> bg_start -> bg_end -> resource_cleanup |
||||
|
assert events == ["resource_start", "bg_start", "bg_end", "resource_cleanup"] |
||||
|
|
||||
|
|
||||
|
# Test 3: Multiple yield dependencies with different cleanup times |
||||
|
def test_multiple_yield_deps_cleanup_order(): |
||||
|
""" |
||||
|
Test that multiple yield dependencies clean up in the correct order. |
||||
|
""" |
||||
|
cleanup_order: list[str] = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def dep_a(): |
||||
|
yield "A" |
||||
|
cleanup_order.append("A") |
||||
|
|
||||
|
async def dep_b(): |
||||
|
yield "B" |
||||
|
cleanup_order.append("B") |
||||
|
|
||||
|
async def dep_c(a: str = Depends(dep_a)): |
||||
|
yield f"C({a})" |
||||
|
cleanup_order.append("C") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
a: str = Depends(dep_a), |
||||
|
b: str = Depends(dep_b), |
||||
|
c: str = Depends(dep_c), |
||||
|
): |
||||
|
async def bg_task(): |
||||
|
cleanup_order.append("BG") |
||||
|
|
||||
|
background_tasks.add_task(bg_task) |
||||
|
return {"a": a, "b": b, "c": c} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
cleanup_order.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Background task should run before cleanups |
||||
|
bg_idx = cleanup_order.index("BG") |
||||
|
# C depends on A, so C cleanup should happen before A |
||||
|
c_idx = cleanup_order.index("C") |
||||
|
a_idx = cleanup_order.index("A") |
||||
|
|
||||
|
assert bg_idx < c_idx, "Background task should run before dependency cleanup" |
||||
|
assert c_idx < a_idx, "C should cleanup before A (dependency order)" |
||||
|
|
||||
|
|
||||
|
# Test 4: Exception in background task should not affect cleanup |
||||
|
def test_exception_in_background_task(): |
||||
|
""" |
||||
|
Test that exceptions in background tasks don't prevent cleanup. |
||||
|
""" |
||||
|
events: list[str] = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_resource(): |
||||
|
events.append("setup") |
||||
|
try: |
||||
|
yield "resource" |
||||
|
finally: |
||||
|
events.append("cleanup") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: str = Depends(get_resource), |
||||
|
): |
||||
|
async def failing_bg_task(): |
||||
|
events.append("bg_task") |
||||
|
raise ValueError("Background task error") |
||||
|
|
||||
|
background_tasks.add_task(failing_bg_task) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
# Test with raise_server_exceptions=True to see the actual exception |
||||
|
client = TestClient(app, raise_server_exceptions=True) |
||||
|
events.clear() |
||||
|
|
||||
|
# Single task failure: original exception is re-raised (backward compatible) |
||||
|
with pytest.raises(ValueError, match="Background task error"): |
||||
|
client.get("/test") |
||||
|
|
||||
|
# Cleanup should still happen even if background task fails |
||||
|
assert "setup" in events |
||||
|
assert "bg_task" in events |
||||
|
assert "cleanup" in events |
||||
|
|
||||
|
|
||||
|
# Test 5: Test with generator (sync) dependencies under async workload |
||||
|
def test_sync_generator_with_async_background_task(): |
||||
|
""" |
||||
|
Test sync generator dependencies with async background tasks. |
||||
|
""" |
||||
|
events: list[str] = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
def sync_resource(): |
||||
|
events.append("sync_setup") |
||||
|
yield "sync_resource" |
||||
|
events.append("sync_cleanup") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: str = Depends(sync_resource), |
||||
|
): |
||||
|
async def async_bg_task(): |
||||
|
events.append("async_bg_start") |
||||
|
await asyncio.sleep(0.01) |
||||
|
events.append("async_bg_end") |
||||
|
|
||||
|
background_tasks.add_task(async_bg_task) |
||||
|
return {"resource": resource} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
events.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Sync resource cleanup should happen after async background task |
||||
|
assert events == [ |
||||
|
"sync_setup", |
||||
|
"async_bg_start", |
||||
|
"async_bg_end", |
||||
|
"sync_cleanup", |
||||
|
] |
||||
|
|
||||
|
|
||||
|
# Test 6: Test with concurrent modification of shared state |
||||
|
def test_concurrent_state_modification(): |
||||
|
""" |
||||
|
Test behavior when background task modifies shared state that |
||||
|
dependency cleanup also accesses. |
||||
|
""" |
||||
|
shared_state = {"value": 0, "history": []} |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_state(): |
||||
|
shared_state["history"].append(f"setup:{shared_state['value']}") |
||||
|
shared_state["value"] = 1 |
||||
|
yield shared_state |
||||
|
shared_state["history"].append(f"cleanup:{shared_state['value']}") |
||||
|
shared_state["value"] = 0 |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
state: dict = Depends(get_state), |
||||
|
): |
||||
|
async def bg_task(s: dict): |
||||
|
s["history"].append(f"bg:{s['value']}") |
||||
|
s["value"] = 2 |
||||
|
|
||||
|
background_tasks.add_task(bg_task, state) |
||||
|
return {"value": state["value"]} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
shared_state["value"] = 0 |
||||
|
shared_state["history"] = [] |
||||
|
|
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# History should show: setup (value 0) -> bg (value 1) -> cleanup (value 2) |
||||
|
assert shared_state["history"] == ["setup:0", "bg:1", "cleanup:2"] |
||||
|
|
||||
|
|
||||
|
# Test 7: Test deeply nested dependencies with background tasks |
||||
|
def test_deeply_nested_dependencies(): |
||||
|
""" |
||||
|
Test deeply nested dependencies (4+ levels) with background tasks. |
||||
|
""" |
||||
|
cleanup_order: list[str] = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def level1(): |
||||
|
yield "L1" |
||||
|
cleanup_order.append("L1") |
||||
|
|
||||
|
async def level2(l1: str = Depends(level1)): |
||||
|
yield f"L2({l1})" |
||||
|
cleanup_order.append("L2") |
||||
|
|
||||
|
async def level3(l2: str = Depends(level2)): |
||||
|
yield f"L3({l2})" |
||||
|
cleanup_order.append("L3") |
||||
|
|
||||
|
async def level4(l3: str = Depends(level3)): |
||||
|
yield f"L4({l3})" |
||||
|
cleanup_order.append("L4") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
l4: str = Depends(level4), |
||||
|
): |
||||
|
async def bg_task(value: str): |
||||
|
cleanup_order.append(f"BG:{value}") |
||||
|
|
||||
|
background_tasks.add_task(bg_task, l4) |
||||
|
return {"value": l4} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
cleanup_order.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
assert response.json()["value"] == "L4(L3(L2(L1)))" |
||||
|
|
||||
|
# BG should run first, then cleanups in reverse order (L4 -> L3 -> L2 -> L1) |
||||
|
assert cleanup_order[0] == "BG:L4(L3(L2(L1)))" |
||||
|
assert cleanup_order[1:] == ["L4", "L3", "L2", "L1"] |
||||
|
|
||||
|
|
||||
|
# Test 8: Test that function-scoped deps with BackgroundTasks logs warning or error |
||||
|
def test_function_scope_logs_potential_issue(): |
||||
|
""" |
||||
|
Document that function scope with background tasks can cause issues. |
||||
|
The resource is cleaned up before the background task runs. |
||||
|
""" |
||||
|
events: list[str] = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_resource(): |
||||
|
events.append("setup") |
||||
|
resource = {"active": True} |
||||
|
yield resource |
||||
|
events.append("cleanup") |
||||
|
resource["active"] = False |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: dict = Depends(get_resource, scope="function"), |
||||
|
): |
||||
|
async def bg_task(res: dict): |
||||
|
events.append(f"bg:active={res['active']}") |
||||
|
|
||||
|
background_tasks.add_task(bg_task, resource) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
events.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# With function scope, cleanup happens BEFORE background task |
||||
|
# This is the documented "inconsistent" behavior |
||||
|
assert events == ["setup", "cleanup", "bg:active=False"] |
||||
|
|
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
pytest.main([__file__, "-v"]) |
||||
@ -0,0 +1,287 @@ |
|||||
|
""" |
||||
|
Tests reproducing actual bugs related to BackgroundTasks and yield dependencies. |
||||
|
|
||||
|
Based on GitHub issues: |
||||
|
1. #14137 - Background tasks added after yield not executed |
||||
|
2. Starlette #2640 - If first background task fails, remaining tasks don't run |
||||
|
3. Starlette #1438 - Background tasks cancelled on client disconnect |
||||
|
""" |
||||
|
|
||||
|
import asyncio |
||||
|
import warnings |
||||
|
from typing import Generator |
||||
|
|
||||
|
import pytest |
||||
|
from fastapi import BackgroundTasks, Depends, FastAPI |
||||
|
from fastapi.testclient import TestClient |
||||
|
|
||||
|
|
||||
|
class TestBackgroundTasksAfterYield: |
||||
|
""" |
||||
|
Bug: Background tasks added after yield are not executed. |
||||
|
|
||||
|
When code after `yield` in a dependency adds a background task, |
||||
|
that task is silently ignored because the response cycle has completed. |
||||
|
""" |
||||
|
|
||||
|
def test_bg_task_added_before_yield_executes(self): |
||||
|
"""Control test: tasks added before yield should execute.""" |
||||
|
executed = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def dependency_with_yield(background_tasks: BackgroundTasks): |
||||
|
background_tasks.add_task(lambda: executed.append("before_yield")) |
||||
|
yield "value" |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(dep: str = Depends(dependency_with_yield)): |
||||
|
return {"dep": dep} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
executed.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
assert "before_yield" in executed |
||||
|
|
||||
|
def test_bg_task_added_after_yield_warns_and_not_executed(self): |
||||
|
""" |
||||
|
FIXED: Tasks added after yield now emit a warning instead of being silently dropped. |
||||
|
|
||||
|
Previously, this was a silent failure. Now users get a clear warning |
||||
|
explaining why the task won't run. |
||||
|
""" |
||||
|
executed = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def dependency_with_yield(background_tasks: BackgroundTasks): |
||||
|
executed.append("before_yield") |
||||
|
yield "value" |
||||
|
# This code runs during cleanup |
||||
|
executed.append("after_yield_code_runs") |
||||
|
# Now this emits a warning instead of silently failing |
||||
|
background_tasks.add_task(lambda: executed.append("after_yield_task")) |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(dep: str = Depends(dependency_with_yield)): |
||||
|
return {"dep": dep} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
executed.clear() |
||||
|
|
||||
|
# Capture warnings during the request |
||||
|
with warnings.catch_warnings(record=True) as w: |
||||
|
warnings.simplefilter("always") |
||||
|
response = client.get("/test") |
||||
|
|
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# The code after yield runs... |
||||
|
assert "after_yield_code_runs" in executed |
||||
|
# ...the background task does NOT execute (as before) |
||||
|
assert "after_yield_task" not in executed |
||||
|
# BUT now we get a warning about it (FIX!) |
||||
|
assert len(w) == 1 |
||||
|
assert "Background task added after tasks have already been executed" in str(w[0].message) |
||||
|
|
||||
|
|
||||
|
class TestBackgroundTaskFailureCascade: |
||||
|
""" |
||||
|
Bug: If first background task fails, remaining tasks don't run. |
||||
|
|
||||
|
From Starlette Discussion #2640: When one background task raises |
||||
|
an exception, subsequent tasks in the queue are skipped. |
||||
|
""" |
||||
|
|
||||
|
def test_first_task_failure_does_not_stop_remaining_tasks(self): |
||||
|
""" |
||||
|
FIXED: Remaining tasks now run even if an earlier task fails. |
||||
|
|
||||
|
Previously, if the first task failed, subsequent tasks wouldn't run. |
||||
|
With the fix, all tasks are attempted regardless of earlier failures. |
||||
|
""" |
||||
|
executed = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(background_tasks: BackgroundTasks): |
||||
|
async def failing_task(): |
||||
|
executed.append("task1_start") |
||||
|
raise ValueError("Task 1 failed") |
||||
|
|
||||
|
async def task2(): |
||||
|
executed.append("task2_executed") |
||||
|
|
||||
|
async def task3(): |
||||
|
executed.append("task3_executed") |
||||
|
|
||||
|
background_tasks.add_task(failing_task) |
||||
|
background_tasks.add_task(task2) |
||||
|
background_tasks.add_task(task3) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=False) |
||||
|
executed.clear() |
||||
|
response = client.get("/test") |
||||
|
|
||||
|
# Response succeeds because error happens in background |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# First task starts |
||||
|
assert "task1_start" in executed |
||||
|
|
||||
|
# FIXED: task2 and task3 now execute even though task1 failed |
||||
|
assert "task2_executed" in executed |
||||
|
assert "task3_executed" in executed |
||||
|
|
||||
|
def test_independent_tasks_all_run(self): |
||||
|
""" |
||||
|
FIXED: Independent tasks now all run regardless of others failing. |
||||
|
|
||||
|
With the fix, task3 runs even though the preceding task failed. |
||||
|
""" |
||||
|
executed = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(background_tasks: BackgroundTasks): |
||||
|
async def task1(): |
||||
|
executed.append("task1") |
||||
|
|
||||
|
async def failing_task(): |
||||
|
executed.append("failing_start") |
||||
|
raise ValueError("I failed") |
||||
|
|
||||
|
async def task3(): |
||||
|
executed.append("task3") |
||||
|
|
||||
|
background_tasks.add_task(task1) |
||||
|
background_tasks.add_task(failing_task) |
||||
|
background_tasks.add_task(task3) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=False) |
||||
|
executed.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Task 1 runs |
||||
|
assert "task1" in executed |
||||
|
# Failing task starts |
||||
|
assert "failing_start" in executed |
||||
|
# FIXED: Task 3 now runs even though the preceding task failed |
||||
|
assert "task3" in executed |
||||
|
|
||||
|
|
||||
|
class TestDependencyCleanupWithFailingTasks: |
||||
|
""" |
||||
|
Test that dependency cleanup happens even when background tasks fail. |
||||
|
""" |
||||
|
|
||||
|
def test_cleanup_happens_despite_bg_task_failure(self): |
||||
|
""" |
||||
|
Verify that yield dependency cleanup runs even if a background task fails. |
||||
|
""" |
||||
|
events = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_resource(): |
||||
|
events.append("setup") |
||||
|
try: |
||||
|
yield "resource" |
||||
|
finally: |
||||
|
events.append("cleanup") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: str = Depends(get_resource), |
||||
|
): |
||||
|
async def failing_task(): |
||||
|
events.append("task_start") |
||||
|
raise ValueError("Task failed") |
||||
|
|
||||
|
background_tasks.add_task(failing_task) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=False) |
||||
|
events.clear() |
||||
|
response = client.get("/test") |
||||
|
|
||||
|
# Response should succeed (bg task error is after response) |
||||
|
assert response.status_code == 200 |
||||
|
assert "setup" in events |
||||
|
assert "task_start" in events |
||||
|
# Cleanup should happen even though task failed |
||||
|
assert "cleanup" in events |
||||
|
|
||||
|
|
||||
|
class TestCleanupOrderConsistency: |
||||
|
""" |
||||
|
Test that cleanup order is consistent and predictable. |
||||
|
""" |
||||
|
|
||||
|
def test_cleanup_order_with_multiple_dependencies(self): |
||||
|
""" |
||||
|
Verify cleanup happens in reverse order of setup. |
||||
|
""" |
||||
|
events = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def dep_a(): |
||||
|
events.append("setup_a") |
||||
|
try: |
||||
|
yield "A" |
||||
|
finally: |
||||
|
events.append("cleanup_a") |
||||
|
|
||||
|
async def dep_b(a: str = Depends(dep_a)): |
||||
|
events.append("setup_b") |
||||
|
try: |
||||
|
yield f"B({a})" |
||||
|
finally: |
||||
|
events.append("cleanup_b") |
||||
|
|
||||
|
async def dep_c(b: str = Depends(dep_b)): |
||||
|
events.append("setup_c") |
||||
|
try: |
||||
|
yield f"C({b})" |
||||
|
finally: |
||||
|
events.append("cleanup_c") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
c: str = Depends(dep_c), |
||||
|
): |
||||
|
async def bg_task(): |
||||
|
events.append("bg_task") |
||||
|
|
||||
|
background_tasks.add_task(bg_task) |
||||
|
return {"c": c} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
events.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Verify setup order: A -> B -> C |
||||
|
assert events.index("setup_a") < events.index("setup_b") |
||||
|
assert events.index("setup_b") < events.index("setup_c") |
||||
|
|
||||
|
# Background task runs after endpoint but before cleanup |
||||
|
assert events.index("bg_task") > events.index("setup_c") |
||||
|
assert events.index("bg_task") < events.index("cleanup_c") |
||||
|
|
||||
|
# Verify cleanup order: C -> B -> A (reverse of setup) |
||||
|
assert events.index("cleanup_c") < events.index("cleanup_b") |
||||
|
assert events.index("cleanup_b") < events.index("cleanup_a") |
||||
|
|
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
pytest.main([__file__, "-v"]) |
||||
@ -0,0 +1,344 @@ |
|||||
|
""" |
||||
|
Test accessing dependencies from within BackgroundTasks. |
||||
|
|
||||
|
This tests scenarios where background tasks try to access dependency-provided |
||||
|
resources directly, which can lead to inconsistent behavior if the dependency |
||||
|
has been cleaned up. |
||||
|
""" |
||||
|
|
||||
|
import asyncio |
||||
|
from typing import Any |
||||
|
|
||||
|
import pytest |
||||
|
from fastapi import BackgroundTasks, Depends, FastAPI, Request |
||||
|
from fastapi.testclient import TestClient |
||||
|
|
||||
|
|
||||
|
# Scenario 1: Background task accessing a resource passed to it |
||||
|
def test_bg_task_with_passed_resource(): |
||||
|
""" |
||||
|
Test that a resource passed to a background task remains valid |
||||
|
when using default (request) scope. |
||||
|
""" |
||||
|
events: list[str] = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
class DatabaseConnection: |
||||
|
def __init__(self): |
||||
|
self.connected = False |
||||
|
self.queries: list[str] = [] |
||||
|
|
||||
|
def connect(self): |
||||
|
self.connected = True |
||||
|
events.append("db_connected") |
||||
|
|
||||
|
def disconnect(self): |
||||
|
self.connected = False |
||||
|
events.append("db_disconnected") |
||||
|
|
||||
|
def query(self, sql: str) -> str: |
||||
|
if not self.connected: |
||||
|
events.append(f"query_failed:{sql}") |
||||
|
raise RuntimeError("Not connected") |
||||
|
events.append(f"query_success:{sql}") |
||||
|
self.queries.append(sql) |
||||
|
return f"result:{sql}" |
||||
|
|
||||
|
async def get_db(): |
||||
|
db = DatabaseConnection() |
||||
|
db.connect() |
||||
|
yield db |
||||
|
db.disconnect() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
db: DatabaseConnection = Depends(get_db), |
||||
|
): |
||||
|
async def bg_query(conn: DatabaseConnection): |
||||
|
# This should work because default scope keeps db alive |
||||
|
result = conn.query("SELECT * FROM bg_table") |
||||
|
events.append(f"bg_got:{result}") |
||||
|
|
||||
|
background_tasks.add_task(bg_query, db) |
||||
|
return {"status": "queued"} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
events.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Verify query succeeded in background task |
||||
|
assert "db_connected" in events |
||||
|
assert "query_success:SELECT * FROM bg_table" in events |
||||
|
assert "bg_got:result:SELECT * FROM bg_table" in events |
||||
|
assert "db_disconnected" in events |
||||
|
|
||||
|
# Verify order: connect -> query -> disconnect |
||||
|
connect_idx = events.index("db_connected") |
||||
|
query_idx = events.index("query_success:SELECT * FROM bg_table") |
||||
|
disconnect_idx = events.index("db_disconnected") |
||||
|
assert connect_idx < query_idx < disconnect_idx |
||||
|
|
||||
|
|
||||
|
# Scenario 2: Background task accessing function-scoped resource (FAILS) |
||||
|
def test_bg_task_with_function_scoped_resource(): |
||||
|
""" |
||||
|
Test that a function-scoped resource is cleaned up before background task runs. |
||||
|
This demonstrates the potential pitfall of using scope="function" with BackgroundTasks. |
||||
|
""" |
||||
|
events: list[str] = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
class DatabaseConnection: |
||||
|
def __init__(self): |
||||
|
self.connected = False |
||||
|
|
||||
|
def connect(self): |
||||
|
self.connected = True |
||||
|
events.append("db_connected") |
||||
|
|
||||
|
def disconnect(self): |
||||
|
self.connected = False |
||||
|
events.append("db_disconnected") |
||||
|
|
||||
|
def query(self, sql: str) -> str: |
||||
|
if not self.connected: |
||||
|
events.append(f"query_failed:{sql}") |
||||
|
return "ERROR: Not connected" |
||||
|
events.append(f"query_success:{sql}") |
||||
|
return f"result:{sql}" |
||||
|
|
||||
|
async def get_db(): |
||||
|
db = DatabaseConnection() |
||||
|
db.connect() |
||||
|
yield db |
||||
|
db.disconnect() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
db: DatabaseConnection = Depends(get_db, scope="function"), |
||||
|
): |
||||
|
async def bg_query(conn: DatabaseConnection): |
||||
|
# This will FAIL because function scope cleans up before bg task |
||||
|
result = conn.query("SELECT * FROM bg_table") |
||||
|
events.append(f"bg_got:{result}") |
||||
|
|
||||
|
background_tasks.add_task(bg_query, db) |
||||
|
return {"status": "queued"} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
events.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# With function scope, db is disconnected BEFORE bg task runs |
||||
|
assert "db_connected" in events |
||||
|
assert "db_disconnected" in events |
||||
|
assert "query_failed:SELECT * FROM bg_table" in events |
||||
|
|
||||
|
# Verify problematic order: connect -> disconnect -> query_failed |
||||
|
connect_idx = events.index("db_connected") |
||||
|
disconnect_idx = events.index("db_disconnected") |
||||
|
query_idx = events.index("query_failed:SELECT * FROM bg_table") |
||||
|
assert connect_idx < disconnect_idx < query_idx |
||||
|
|
||||
|
|
||||
|
# Scenario 3: Multiple concurrent requests with shared mutable state |
||||
|
def test_concurrent_requests_shared_state(): |
||||
|
""" |
||||
|
Test behavior when multiple concurrent requests share mutable state |
||||
|
through a dependency. |
||||
|
""" |
||||
|
events: list[str] = [] |
||||
|
request_counter = {"value": 0} |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_request_id(): |
||||
|
request_counter["value"] += 1 |
||||
|
request_id = request_counter["value"] |
||||
|
events.append(f"setup:{request_id}") |
||||
|
yield request_id |
||||
|
events.append(f"cleanup:{request_id}") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
request_id: int = Depends(get_request_id), |
||||
|
): |
||||
|
async def bg_task(rid: int): |
||||
|
await asyncio.sleep(0.01) # Small delay |
||||
|
events.append(f"bg:{rid}") |
||||
|
|
||||
|
background_tasks.add_task(bg_task, request_id) |
||||
|
return {"request_id": request_id} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
events.clear() |
||||
|
request_counter["value"] = 0 |
||||
|
|
||||
|
# Make multiple sequential requests |
||||
|
for _ in range(3): |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Verify each request has proper setup -> bg -> cleanup sequence |
||||
|
# (Note: with default scope, cleanup happens after bg task) |
||||
|
assert len([e for e in events if e.startswith("setup:")]) == 3 |
||||
|
assert len([e for e in events if e.startswith("bg:")]) == 3 |
||||
|
assert len([e for e in events if e.startswith("cleanup:")]) == 3 |
||||
|
|
||||
|
|
||||
|
# Scenario 4: Nested async generators with background tasks |
||||
|
def test_nested_async_generators_with_bg_tasks(): |
||||
|
""" |
||||
|
Test deeply nested async generators and their interaction with background tasks. |
||||
|
""" |
||||
|
events: list[str] = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def outer(): |
||||
|
events.append("outer_start") |
||||
|
yield "outer" |
||||
|
events.append("outer_end") |
||||
|
|
||||
|
async def middle(outer_val: str = Depends(outer)): |
||||
|
events.append(f"middle_start:{outer_val}") |
||||
|
yield f"middle({outer_val})" |
||||
|
events.append("middle_end") |
||||
|
|
||||
|
async def inner(middle_val: str = Depends(middle)): |
||||
|
events.append(f"inner_start:{middle_val}") |
||||
|
yield f"inner({middle_val})" |
||||
|
events.append("inner_end") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
value: str = Depends(inner), |
||||
|
): |
||||
|
async def bg_task(val: str): |
||||
|
events.append(f"bg:{val}") |
||||
|
|
||||
|
background_tasks.add_task(bg_task, value) |
||||
|
return {"value": value} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
events.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
assert response.json()["value"] == "inner(middle(outer))" |
||||
|
|
||||
|
# Verify setup order: outer -> middle -> inner |
||||
|
assert events.index("outer_start") < events.index("middle_start:outer") |
||||
|
assert events.index("middle_start:outer") < events.index("inner_start:middle(outer)") |
||||
|
|
||||
|
# Verify bg task runs before cleanup |
||||
|
bg_idx = events.index("bg:inner(middle(outer))") |
||||
|
assert bg_idx < events.index("inner_end") |
||||
|
|
||||
|
# Verify cleanup order (reverse): inner -> middle -> outer |
||||
|
assert events.index("inner_end") < events.index("middle_end") |
||||
|
assert events.index("middle_end") < events.index("outer_end") |
||||
|
|
||||
|
|
||||
|
# Scenario 5: Background task that adds more background tasks |
||||
|
def test_bg_task_adding_more_bg_tasks(): |
||||
|
""" |
||||
|
Test a background task that adds additional background tasks. |
||||
|
""" |
||||
|
events: list[str] = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_resource(): |
||||
|
events.append("resource_setup") |
||||
|
yield {"active": True} |
||||
|
events.append("resource_cleanup") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
request: Request, |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: dict = Depends(get_resource), |
||||
|
): |
||||
|
async def first_bg_task(res: dict, tasks: BackgroundTasks): |
||||
|
events.append(f"first_bg:active={res['active']}") |
||||
|
# Note: Adding tasks from within a task might not work as expected |
||||
|
# because the BackgroundTasks object is tied to the response |
||||
|
|
||||
|
async def second_bg_task(res: dict): |
||||
|
events.append(f"second_bg:active={res['active']}") |
||||
|
|
||||
|
background_tasks.add_task(first_bg_task, resource, background_tasks) |
||||
|
background_tasks.add_task(second_bg_task, resource) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
events.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Both background tasks should see active=True (default scope) |
||||
|
assert "first_bg:active=True" in events |
||||
|
assert "second_bg:active=True" in events |
||||
|
assert "resource_cleanup" in events |
||||
|
|
||||
|
|
||||
|
# Scenario 6: Error in one background task doesn't affect others |
||||
|
def test_error_in_one_bg_task_others_continue(): |
||||
|
""" |
||||
|
Test that an error in one background task doesn't prevent cleanup or other tasks. |
||||
|
Note: Cleanup code MUST be in a finally block to run when exceptions occur. |
||||
|
|
||||
|
With the fix, all tasks run even if one fails. |
||||
|
""" |
||||
|
events: list[str] = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_resource(): |
||||
|
events.append("setup") |
||||
|
try: |
||||
|
yield "resource" |
||||
|
finally: |
||||
|
events.append("cleanup") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: str = Depends(get_resource), |
||||
|
): |
||||
|
async def failing_task(): |
||||
|
events.append("failing_start") |
||||
|
raise ValueError("Task failed") |
||||
|
|
||||
|
async def succeeding_task(): |
||||
|
events.append("succeeding") |
||||
|
|
||||
|
background_tasks.add_task(failing_task) |
||||
|
background_tasks.add_task(succeeding_task) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=True) |
||||
|
events.clear() |
||||
|
|
||||
|
# Single task failure: original exception re-raised (backward compatible) |
||||
|
with pytest.raises(ValueError, match="Task failed"): |
||||
|
client.get("/test") |
||||
|
|
||||
|
# Cleanup should still happen |
||||
|
assert "setup" in events |
||||
|
assert "failing_start" in events |
||||
|
assert "cleanup" in events |
||||
|
# FIXED: Second task now runs even though first task failed |
||||
|
assert "succeeding" in events |
||||
|
|
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
pytest.main([__file__, "-v"]) |
||||
@ -0,0 +1,411 @@ |
|||||
|
""" |
||||
|
Tests verifying the fixes for BackgroundTasks bugs. |
||||
|
|
||||
|
Fixed bugs: |
||||
|
1. First task failure now doesn't stop remaining tasks |
||||
|
2. Tasks added after execution now emit a warning |
||||
|
""" |
||||
|
|
||||
|
import warnings |
||||
|
from typing import Generator |
||||
|
|
||||
|
import pytest |
||||
|
from fastapi import BackgroundTasks, Depends, FastAPI |
||||
|
from fastapi.background import BackgroundTaskError |
||||
|
from fastapi.testclient import TestClient |
||||
|
|
||||
|
|
||||
|
class TestBackgroundTaskErrorIsolation: |
||||
|
""" |
||||
|
Test that task failures don't stop remaining tasks. |
||||
|
""" |
||||
|
|
||||
|
def test_all_tasks_run_despite_failure(self): |
||||
|
""" |
||||
|
FIX: All tasks should run even if one fails in the middle. |
||||
|
""" |
||||
|
executed = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(background_tasks: BackgroundTasks): |
||||
|
async def task1(): |
||||
|
executed.append("task1") |
||||
|
|
||||
|
async def failing_task(): |
||||
|
executed.append("failing_start") |
||||
|
raise ValueError("I failed") |
||||
|
|
||||
|
async def task3(): |
||||
|
executed.append("task3") |
||||
|
|
||||
|
background_tasks.add_task(task1) |
||||
|
background_tasks.add_task(failing_task) |
||||
|
background_tasks.add_task(task3) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=False) |
||||
|
executed.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# FIX VERIFICATION: All three tasks should have run |
||||
|
assert "task1" in executed |
||||
|
assert "failing_start" in executed |
||||
|
assert "task3" in executed # This now runs! |
||||
|
|
||||
|
def test_multiple_failures_all_logged(self): |
||||
|
""" |
||||
|
Multiple task failures should all be captured. |
||||
|
""" |
||||
|
executed = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(background_tasks: BackgroundTasks): |
||||
|
async def failing1(): |
||||
|
executed.append("failing1") |
||||
|
raise ValueError("Error 1") |
||||
|
|
||||
|
async def failing2(): |
||||
|
executed.append("failing2") |
||||
|
raise RuntimeError("Error 2") |
||||
|
|
||||
|
async def success(): |
||||
|
executed.append("success") |
||||
|
|
||||
|
background_tasks.add_task(failing1) |
||||
|
background_tasks.add_task(success) |
||||
|
background_tasks.add_task(failing2) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=False) |
||||
|
executed.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# All tasks should have executed |
||||
|
assert "failing1" in executed |
||||
|
assert "success" in executed |
||||
|
assert "failing2" in executed |
||||
|
|
||||
|
def test_single_task_failure_raises_original_exception(self): |
||||
|
""" |
||||
|
Single task failure should raise the original exception (backward compatible). |
||||
|
""" |
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(background_tasks: BackgroundTasks): |
||||
|
async def failing(): |
||||
|
raise ValueError("Task failed") |
||||
|
|
||||
|
background_tasks.add_task(failing) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=True) |
||||
|
|
||||
|
# Single failure: original exception raised for backward compatibility |
||||
|
with pytest.raises(ValueError, match="Task failed"): |
||||
|
client.get("/test") |
||||
|
|
||||
|
def test_multiple_task_failures_raise_background_task_error(self): |
||||
|
""" |
||||
|
Multiple task failures should raise BackgroundTaskError with all errors. |
||||
|
""" |
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(background_tasks: BackgroundTasks): |
||||
|
async def failing1(): |
||||
|
raise ValueError("Error 1") |
||||
|
|
||||
|
async def failing2(): |
||||
|
raise RuntimeError("Error 2") |
||||
|
|
||||
|
background_tasks.add_task(failing1) |
||||
|
background_tasks.add_task(failing2) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=True) |
||||
|
|
||||
|
# Multiple failures: BackgroundTaskError raised |
||||
|
with pytest.raises(BackgroundTaskError) as exc_info: |
||||
|
client.get("/test") |
||||
|
|
||||
|
assert len(exc_info.value.errors) == 2 |
||||
|
assert isinstance(exc_info.value.errors[0][1], ValueError) |
||||
|
assert isinstance(exc_info.value.errors[1][1], RuntimeError) |
||||
|
|
||||
|
|
||||
|
class TestBackgroundTaskAfterYieldWarning: |
||||
|
""" |
||||
|
Test that adding tasks after execution emits a warning. |
||||
|
""" |
||||
|
|
||||
|
def test_warning_when_task_added_after_execution(self): |
||||
|
""" |
||||
|
FIX: Warning should be emitted when task is added after execution. |
||||
|
""" |
||||
|
executed = [] |
||||
|
warnings_captured = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def dependency_with_yield(background_tasks: BackgroundTasks): |
||||
|
background_tasks.add_task(lambda: executed.append("before_yield")) |
||||
|
yield "value" |
||||
|
# Capture the warning |
||||
|
with warnings.catch_warnings(record=True) as w: |
||||
|
warnings.simplefilter("always") |
||||
|
background_tasks.add_task(lambda: executed.append("after_yield")) |
||||
|
if w: |
||||
|
warnings_captured.extend(w) |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(dep: str = Depends(dependency_with_yield)): |
||||
|
return {"dep": dep} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
executed.clear() |
||||
|
warnings_captured.clear() |
||||
|
|
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Task added before yield should have executed |
||||
|
assert "before_yield" in executed |
||||
|
|
||||
|
# Task added after yield should NOT have executed |
||||
|
assert "after_yield" not in executed |
||||
|
|
||||
|
# Warning should have been captured |
||||
|
assert len(warnings_captured) == 1 |
||||
|
assert "Background task added after tasks have already been executed" in str( |
||||
|
warnings_captured[0].message |
||||
|
) |
||||
|
|
||||
|
def test_no_warning_for_normal_task_addition(self): |
||||
|
""" |
||||
|
No warning should be emitted for normal task addition. |
||||
|
""" |
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(background_tasks: BackgroundTasks): |
||||
|
with warnings.catch_warnings(record=True) as w: |
||||
|
warnings.simplefilter("always") |
||||
|
background_tasks.add_task(lambda: None) |
||||
|
background_tasks.add_task(lambda: None) |
||||
|
assert len(w) == 0 # No warnings |
||||
|
|
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
|
||||
|
class TestDependencyCleanupStillWorks: |
||||
|
""" |
||||
|
Verify that dependency cleanup still works correctly with the fixes. |
||||
|
""" |
||||
|
|
||||
|
def test_cleanup_runs_after_all_tasks(self): |
||||
|
""" |
||||
|
Dependency cleanup should run after all background tasks complete. |
||||
|
""" |
||||
|
events = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_resource(): |
||||
|
events.append("setup") |
||||
|
try: |
||||
|
yield "resource" |
||||
|
finally: |
||||
|
events.append("cleanup") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: str = Depends(get_resource), |
||||
|
): |
||||
|
async def task1(): |
||||
|
events.append("task1") |
||||
|
|
||||
|
async def task2(): |
||||
|
events.append("task2") |
||||
|
|
||||
|
background_tasks.add_task(task1) |
||||
|
background_tasks.add_task(task2) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
events.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Verify order: setup -> tasks -> cleanup |
||||
|
assert events.index("setup") < events.index("task1") |
||||
|
assert events.index("task1") < events.index("task2") |
||||
|
assert events.index("task2") < events.index("cleanup") |
||||
|
|
||||
|
def test_cleanup_runs_even_when_tasks_fail(self): |
||||
|
""" |
||||
|
Dependency cleanup should run even if background tasks fail. |
||||
|
""" |
||||
|
events = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def get_resource(): |
||||
|
events.append("setup") |
||||
|
try: |
||||
|
yield "resource" |
||||
|
finally: |
||||
|
events.append("cleanup") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: str = Depends(get_resource), |
||||
|
): |
||||
|
async def failing_task(): |
||||
|
events.append("task_start") |
||||
|
raise ValueError("Failed") |
||||
|
|
||||
|
background_tasks.add_task(failing_task) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=False) |
||||
|
events.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Cleanup must happen |
||||
|
assert "setup" in events |
||||
|
assert "task_start" in events |
||||
|
assert "cleanup" in events |
||||
|
|
||||
|
|
||||
|
class TestEdgeCases: |
||||
|
"""Test edge cases and robustness fixes.""" |
||||
|
|
||||
|
def test_task_with_no_name_attribute(self): |
||||
|
"""Test that tasks without __name__ are handled gracefully.""" |
||||
|
executed = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(background_tasks: BackgroundTasks): |
||||
|
# Lambda has __name__ = '<lambda>' |
||||
|
background_tasks.add_task(lambda: executed.append("lambda")) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
executed.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
assert "lambda" in executed |
||||
|
|
||||
|
def test_task_mutation_during_execution(self): |
||||
|
"""Test that task list mutation during execution doesn't cause issues.""" |
||||
|
executed = [] |
||||
|
tasks_ref = {"bg": None} |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(background_tasks: BackgroundTasks): |
||||
|
tasks_ref["bg"] = background_tasks |
||||
|
|
||||
|
async def mutating_task(): |
||||
|
executed.append("task1") |
||||
|
# Try to mutate the tasks list during execution |
||||
|
# This should not affect iteration due to snapshot |
||||
|
tasks_ref["bg"].tasks.append( |
||||
|
BackgroundTasks() # Add garbage |
||||
|
) |
||||
|
|
||||
|
async def task2(): |
||||
|
executed.append("task2") |
||||
|
|
||||
|
background_tasks.add_task(mutating_task) |
||||
|
background_tasks.add_task(task2) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=False) |
||||
|
executed.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Both tasks should have executed despite mutation attempt |
||||
|
assert "task1" in executed |
||||
|
assert "task2" in executed |
||||
|
|
||||
|
@pytest.mark.anyio |
||||
|
async def test_keyboard_interrupt_not_suppressed(self): |
||||
|
"""Test that KeyboardInterrupt is not suppressed.""" |
||||
|
from fastapi.background import BackgroundTasks as BgTasks |
||||
|
|
||||
|
async def raising_task(): |
||||
|
raise KeyboardInterrupt() |
||||
|
|
||||
|
bg = BgTasks() |
||||
|
bg.add_task(raising_task) |
||||
|
|
||||
|
with pytest.raises(KeyboardInterrupt): |
||||
|
await bg() |
||||
|
|
||||
|
@pytest.mark.anyio |
||||
|
async def test_system_exit_not_suppressed(self): |
||||
|
"""Test that SystemExit is not suppressed.""" |
||||
|
from fastapi.background import BackgroundTasks as BgTasks |
||||
|
|
||||
|
async def raising_task(): |
||||
|
raise SystemExit(1) |
||||
|
|
||||
|
bg = BgTasks() |
||||
|
bg.add_task(raising_task) |
||||
|
|
||||
|
with pytest.raises(SystemExit): |
||||
|
await bg() |
||||
|
|
||||
|
def test_base_exception_subclass_caught(self): |
||||
|
"""Test that BaseException subclasses (except critical ones) are caught.""" |
||||
|
executed = [] |
||||
|
|
||||
|
app = FastAPI() |
||||
|
|
||||
|
class CustomBaseException(BaseException): |
||||
|
pass |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint(background_tasks: BackgroundTasks): |
||||
|
async def failing_task(): |
||||
|
executed.append("failing") |
||||
|
raise CustomBaseException("custom error") |
||||
|
|
||||
|
async def second_task(): |
||||
|
executed.append("second") |
||||
|
|
||||
|
background_tasks.add_task(failing_task) |
||||
|
background_tasks.add_task(second_task) |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=False) |
||||
|
executed.clear() |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Both tasks should execute - CustomBaseException should be caught |
||||
|
assert "failing" in executed |
||||
|
assert "second" in executed |
||||
|
|
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
pytest.main([__file__, "-v"]) |
||||
@ -0,0 +1,355 @@ |
|||||
|
""" |
||||
|
Test for async dependencies with cleanup (yield) and BackgroundTasks interaction. |
||||
|
|
||||
|
This tests the scenario where dependencies with cleanup logic (yield) may behave |
||||
|
inconsistently when used inside BackgroundTasks under different scope settings. |
||||
|
|
||||
|
The key issue: |
||||
|
- Dependencies with scope="function" have cleanup run BEFORE BackgroundTasks execute |
||||
|
- Dependencies with scope="request" (default) have cleanup run AFTER BackgroundTasks execute |
||||
|
|
||||
|
This can lead to inconsistent behavior where a resource is cleaned up before |
||||
|
a background task that depends on it has a chance to run. |
||||
|
""" |
||||
|
|
||||
|
import asyncio |
||||
|
from typing import Any |
||||
|
|
||||
|
import pytest |
||||
|
from fastapi import BackgroundTasks, Depends, FastAPI |
||||
|
from fastapi.testclient import TestClient |
||||
|
|
||||
|
|
||||
|
# Track cleanup and background task execution order |
||||
|
execution_log: list[str] = [] |
||||
|
|
||||
|
|
||||
|
def reset_log(): |
||||
|
execution_log.clear() |
||||
|
|
||||
|
|
||||
|
# Test 1: Default scope (request) - cleanup should happen AFTER background task |
||||
|
def test_background_task_with_request_scope_yield_dependency(): |
||||
|
""" |
||||
|
With default (request) scope, dependency cleanup should happen AFTER |
||||
|
the background task completes. |
||||
|
""" |
||||
|
reset_log() |
||||
|
app = FastAPI() |
||||
|
resource_value = {"status": "initialized"} |
||||
|
|
||||
|
async def get_resource(): |
||||
|
execution_log.append("resource_setup") |
||||
|
resource_value["status"] = "active" |
||||
|
yield resource_value |
||||
|
execution_log.append("resource_cleanup") |
||||
|
resource_value["status"] = "cleaned_up" |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: dict = Depends(get_resource), |
||||
|
): |
||||
|
async def bg_task(res: dict): |
||||
|
execution_log.append(f"bg_task_start:status={res['status']}") |
||||
|
await asyncio.sleep(0.01) # Simulate async work |
||||
|
execution_log.append(f"bg_task_end:status={res['status']}") |
||||
|
|
||||
|
background_tasks.add_task(bg_task, resource) |
||||
|
return {"status": resource["status"]} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
assert response.json() == {"status": "active"} |
||||
|
|
||||
|
# Verify order: setup -> bg_task -> cleanup |
||||
|
assert execution_log == [ |
||||
|
"resource_setup", |
||||
|
"bg_task_start:status=active", |
||||
|
"bg_task_end:status=active", |
||||
|
"resource_cleanup", |
||||
|
] |
||||
|
|
||||
|
|
||||
|
# Test 2: Function scope - cleanup happens BEFORE background task (BUG!) |
||||
|
def test_background_task_with_function_scope_yield_dependency(): |
||||
|
""" |
||||
|
With scope="function", dependency cleanup happens BEFORE the background task, |
||||
|
which is inconsistent and potentially problematic. |
||||
|
|
||||
|
This test demonstrates the issue: the resource is cleaned up before |
||||
|
the background task runs, so the background task sees the cleaned up state. |
||||
|
""" |
||||
|
reset_log() |
||||
|
app = FastAPI() |
||||
|
resource_value = {"status": "initialized"} |
||||
|
|
||||
|
async def get_resource(): |
||||
|
execution_log.append("resource_setup") |
||||
|
resource_value["status"] = "active" |
||||
|
yield resource_value |
||||
|
execution_log.append("resource_cleanup") |
||||
|
resource_value["status"] = "cleaned_up" |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
resource: dict = Depends(get_resource, scope="function"), |
||||
|
): |
||||
|
async def bg_task(res: dict): |
||||
|
execution_log.append(f"bg_task_start:status={res['status']}") |
||||
|
await asyncio.sleep(0.01) |
||||
|
execution_log.append(f"bg_task_end:status={res['status']}") |
||||
|
|
||||
|
background_tasks.add_task(bg_task, resource) |
||||
|
return {"status": resource["status"]} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
assert response.json() == {"status": "active"} |
||||
|
|
||||
|
# With function scope, cleanup happens BEFORE background task runs |
||||
|
# This is the inconsistent behavior! |
||||
|
assert execution_log == [ |
||||
|
"resource_setup", |
||||
|
"resource_cleanup", # Cleanup happens first! |
||||
|
"bg_task_start:status=cleaned_up", # BG task sees cleaned up state |
||||
|
"bg_task_end:status=cleaned_up", |
||||
|
] |
||||
|
|
||||
|
|
||||
|
# Test 3: Database-like connection scenario with function scope |
||||
|
def test_database_connection_closed_before_background_task(): |
||||
|
""" |
||||
|
Simulates a real-world scenario where a database connection is closed |
||||
|
before a background task that needs it runs. |
||||
|
""" |
||||
|
reset_log() |
||||
|
app = FastAPI() |
||||
|
|
||||
|
class FakeDBConnection: |
||||
|
def __init__(self): |
||||
|
self.is_open = False |
||||
|
self.data = [] |
||||
|
|
||||
|
def open(self): |
||||
|
self.is_open = True |
||||
|
execution_log.append("db_opened") |
||||
|
|
||||
|
def close(self): |
||||
|
self.is_open = False |
||||
|
execution_log.append("db_closed") |
||||
|
|
||||
|
def query(self): |
||||
|
if not self.is_open: |
||||
|
execution_log.append("db_query_failed:connection_closed") |
||||
|
raise RuntimeError("Connection closed") |
||||
|
execution_log.append("db_query_success") |
||||
|
return {"result": "data"} |
||||
|
|
||||
|
db = FakeDBConnection() |
||||
|
|
||||
|
async def get_db(): |
||||
|
db.open() |
||||
|
yield db |
||||
|
db.close() |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
conn: FakeDBConnection = Depends(get_db, scope="function"), |
||||
|
): |
||||
|
async def bg_task(connection: FakeDBConnection): |
||||
|
try: |
||||
|
result = connection.query() |
||||
|
execution_log.append(f"bg_task_got_result") |
||||
|
except RuntimeError as e: |
||||
|
execution_log.append(f"bg_task_error:{e}") |
||||
|
|
||||
|
background_tasks.add_task(bg_task, conn) |
||||
|
return {"status": "task_queued"} |
||||
|
|
||||
|
client = TestClient(app, raise_server_exceptions=False) |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# The background task fails because connection was closed before it ran |
||||
|
assert "db_closed" in execution_log |
||||
|
assert execution_log.index("db_closed") < execution_log.index("db_query_failed:connection_closed") |
||||
|
|
||||
|
|
||||
|
# Test 4: Nested dependencies with mixed scopes |
||||
|
def test_nested_dependencies_mixed_scopes(): |
||||
|
""" |
||||
|
Test nested dependencies where outer has request scope and inner has function scope. |
||||
|
""" |
||||
|
reset_log() |
||||
|
app = FastAPI() |
||||
|
|
||||
|
async def outer_resource(): |
||||
|
execution_log.append("outer_setup") |
||||
|
yield "outer_active" |
||||
|
execution_log.append("outer_cleanup") |
||||
|
|
||||
|
async def inner_resource(outer: str = Depends(outer_resource)): |
||||
|
execution_log.append(f"inner_setup:outer={outer}") |
||||
|
yield f"inner_active:outer={outer}" |
||||
|
execution_log.append("inner_cleanup") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
inner: str = Depends(inner_resource, scope="function"), |
||||
|
): |
||||
|
async def bg_task(value: str): |
||||
|
execution_log.append(f"bg_task:{value}") |
||||
|
|
||||
|
background_tasks.add_task(bg_task, inner) |
||||
|
return {"value": inner} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Check execution order - inner cleanup should happen before outer |
||||
|
# and before background task due to function scope |
||||
|
print("Execution log:", execution_log) |
||||
|
assert "inner_cleanup" in execution_log |
||||
|
assert "outer_cleanup" in execution_log |
||||
|
|
||||
|
|
||||
|
# Test 5: Multiple background tasks with function scope dependency |
||||
|
def test_multiple_background_tasks_function_scope(): |
||||
|
""" |
||||
|
Multiple background tasks all trying to use a function-scoped dependency. |
||||
|
""" |
||||
|
reset_log() |
||||
|
app = FastAPI() |
||||
|
counter = {"value": 0} |
||||
|
|
||||
|
async def get_counter(): |
||||
|
counter["value"] = 1 |
||||
|
execution_log.append("counter_setup") |
||||
|
yield counter |
||||
|
execution_log.append("counter_cleanup") |
||||
|
counter["value"] = 0 |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
cnt: dict = Depends(get_counter, scope="function"), |
||||
|
): |
||||
|
async def increment(c: dict): |
||||
|
execution_log.append(f"increment:value={c['value']}") |
||||
|
c["value"] += 1 |
||||
|
|
||||
|
background_tasks.add_task(increment, cnt) |
||||
|
background_tasks.add_task(increment, cnt) |
||||
|
background_tasks.add_task(increment, cnt) |
||||
|
return {"initial_value": cnt["value"]} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
assert response.json() == {"initial_value": 1} |
||||
|
|
||||
|
# Check that cleanup happened before background tasks ran |
||||
|
cleanup_idx = execution_log.index("counter_cleanup") |
||||
|
increment_indices = [ |
||||
|
i for i, log in enumerate(execution_log) if log.startswith("increment") |
||||
|
] |
||||
|
|
||||
|
# With function scope, cleanup happens before all increments |
||||
|
assert all(cleanup_idx < idx for idx in increment_indices) |
||||
|
|
||||
|
|
||||
|
# Test 6: Async workload with concurrent requests |
||||
|
def test_concurrent_requests_with_yield_dependencies(): |
||||
|
""" |
||||
|
Test behavior under concurrent async workload. |
||||
|
""" |
||||
|
reset_log() |
||||
|
app = FastAPI() |
||||
|
shared_state = {"active_connections": 0, "max_connections": 0} |
||||
|
|
||||
|
async def get_connection(): |
||||
|
shared_state["active_connections"] += 1 |
||||
|
shared_state["max_connections"] = max( |
||||
|
shared_state["max_connections"], shared_state["active_connections"] |
||||
|
) |
||||
|
execution_log.append(f"conn_open:active={shared_state['active_connections']}") |
||||
|
yield shared_state["active_connections"] |
||||
|
shared_state["active_connections"] -= 1 |
||||
|
execution_log.append(f"conn_close:active={shared_state['active_connections']}") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
conn_id: int = Depends(get_connection, scope="function"), |
||||
|
): |
||||
|
async def bg_task(id: int): |
||||
|
await asyncio.sleep(0.01) |
||||
|
execution_log.append(f"bg_task:conn_id={id}:active={shared_state['active_connections']}") |
||||
|
|
||||
|
background_tasks.add_task(bg_task, conn_id) |
||||
|
return {"conn_id": conn_id} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
|
||||
|
# Make a single request first |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Verify that connection was closed before background task ran |
||||
|
# (due to function scope) |
||||
|
conn_close_indices = [ |
||||
|
i for i, log in enumerate(execution_log) if log.startswith("conn_close") |
||||
|
] |
||||
|
bg_task_indices = [ |
||||
|
i for i, log in enumerate(execution_log) if log.startswith("bg_task") |
||||
|
] |
||||
|
|
||||
|
if conn_close_indices and bg_task_indices: |
||||
|
# Function scope: close happens before bg_task |
||||
|
assert conn_close_indices[0] < bg_task_indices[0] |
||||
|
|
||||
|
|
||||
|
# Test 7: Verify request scope works correctly (control test) |
||||
|
def test_request_scope_preserves_resource_for_background_task(): |
||||
|
""" |
||||
|
Control test: request scope should keep resource active during background task. |
||||
|
""" |
||||
|
reset_log() |
||||
|
app = FastAPI() |
||||
|
resource = {"is_active": False} |
||||
|
|
||||
|
async def get_resource(): |
||||
|
resource["is_active"] = True |
||||
|
execution_log.append("resource_activated") |
||||
|
yield resource |
||||
|
resource["is_active"] = False |
||||
|
execution_log.append("resource_deactivated") |
||||
|
|
||||
|
@app.get("/test") |
||||
|
async def endpoint( |
||||
|
background_tasks: BackgroundTasks, |
||||
|
res: dict = Depends(get_resource), # Default request scope |
||||
|
): |
||||
|
async def bg_task(r: dict): |
||||
|
execution_log.append(f"bg_task:is_active={r['is_active']}") |
||||
|
|
||||
|
background_tasks.add_task(bg_task, res) |
||||
|
return {"is_active": res["is_active"]} |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
response = client.get("/test") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
# Background task should see is_active=True because cleanup happens after |
||||
|
assert "bg_task:is_active=True" in execution_log |
||||
|
|
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
pytest.main([__file__, "-v"]) |
||||
@ -0,0 +1,377 @@ |
|||||
|
""" |
||||
|
Tests for unified lifespan handling. |
||||
|
""" |
||||
|
|
||||
|
import warnings |
||||
|
from collections.abc import AsyncGenerator |
||||
|
from contextlib import asynccontextmanager |
||||
|
from typing import Any |
||||
|
|
||||
|
import pytest |
||||
|
from fastapi import FastAPI, Request |
||||
|
from fastapi.exceptions import FastAPIDeprecationWarning |
||||
|
from fastapi.lifespan import ( |
||||
|
LifespanError, |
||||
|
UnifiedLifespanManager, |
||||
|
create_unified_lifespan, |
||||
|
merge_lifespan_state, |
||||
|
) |
||||
|
from fastapi.testclient import TestClient |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# Test Fixtures |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
@pytest.fixture |
||||
|
def startup_shutdown_log(): |
||||
|
"""Fixture to track startup/shutdown execution order.""" |
||||
|
return [] |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# UnifiedLifespanManager Tests |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class TestUnifiedLifespanManager: |
||||
|
"""Tests for UnifiedLifespanManager class.""" |
||||
|
|
||||
|
@pytest.mark.anyio |
||||
|
async def test_basic_lifespan_execution(self, startup_shutdown_log): |
||||
|
"""Test basic lifespan execution order.""" |
||||
|
|
||||
|
@asynccontextmanager |
||||
|
async def lifespan(app): |
||||
|
startup_shutdown_log.append("startup") |
||||
|
yield {"key": "value"} |
||||
|
startup_shutdown_log.append("shutdown") |
||||
|
|
||||
|
manager = UnifiedLifespanManager(lifespan=lifespan, emit_deprecation_warnings=False) |
||||
|
|
||||
|
with warnings.catch_warnings(): |
||||
|
warnings.simplefilter("ignore", FastAPIDeprecationWarning) |
||||
|
async with manager.lifespan(None) as state: |
||||
|
assert state == {"key": "value"} |
||||
|
startup_shutdown_log.append("running") |
||||
|
|
||||
|
assert startup_shutdown_log == ["startup", "running", "shutdown"] |
||||
|
|
||||
|
@pytest.mark.anyio |
||||
|
async def test_on_event_handlers_execution(self, startup_shutdown_log): |
||||
|
"""Test on_event handlers execute correctly.""" |
||||
|
|
||||
|
async def startup1(): |
||||
|
startup_shutdown_log.append("startup1") |
||||
|
|
||||
|
async def startup2(): |
||||
|
startup_shutdown_log.append("startup2") |
||||
|
|
||||
|
async def shutdown1(): |
||||
|
startup_shutdown_log.append("shutdown1") |
||||
|
|
||||
|
async def shutdown2(): |
||||
|
startup_shutdown_log.append("shutdown2") |
||||
|
|
||||
|
manager = UnifiedLifespanManager( |
||||
|
on_startup=[startup1, startup2], |
||||
|
on_shutdown=[shutdown1, shutdown2], |
||||
|
emit_deprecation_warnings=False, |
||||
|
) |
||||
|
|
||||
|
async with manager.lifespan(None) as state: |
||||
|
startup_shutdown_log.append("running") |
||||
|
|
||||
|
# Startup in order, shutdown in reverse order |
||||
|
assert startup_shutdown_log == [ |
||||
|
"startup1", |
||||
|
"startup2", |
||||
|
"running", |
||||
|
"shutdown2", |
||||
|
"shutdown1", |
||||
|
] |
||||
|
|
||||
|
@pytest.mark.anyio |
||||
|
async def test_startup_failure_rollback(self, startup_shutdown_log): |
||||
|
"""Test rollback when startup fails.""" |
||||
|
|
||||
|
async def startup1(): |
||||
|
startup_shutdown_log.append("startup1") |
||||
|
|
||||
|
async def startup2(): |
||||
|
startup_shutdown_log.append("startup2_failed") |
||||
|
raise ValueError("Startup failed") |
||||
|
|
||||
|
async def shutdown1(): |
||||
|
startup_shutdown_log.append("shutdown1") |
||||
|
|
||||
|
async def shutdown2(): |
||||
|
startup_shutdown_log.append("shutdown2") |
||||
|
|
||||
|
manager = UnifiedLifespanManager( |
||||
|
on_startup=[startup1, startup2], |
||||
|
on_shutdown=[shutdown1, shutdown2], |
||||
|
emit_deprecation_warnings=False, |
||||
|
) |
||||
|
|
||||
|
with pytest.raises(LifespanError) as exc_info: |
||||
|
async with manager.lifespan(None): |
||||
|
startup_shutdown_log.append("should_not_run") |
||||
|
|
||||
|
# Rollback should clean up startup1 only (shutdown1) |
||||
|
assert "startup1" in startup_shutdown_log |
||||
|
assert "startup2_failed" in startup_shutdown_log |
||||
|
assert "should_not_run" not in startup_shutdown_log |
||||
|
# Rollback runs cleanup for completed handlers |
||||
|
assert "shutdown1" in startup_shutdown_log |
||||
|
|
||||
|
# Check error details |
||||
|
assert exc_info.value.phase == "startup" |
||||
|
assert isinstance(exc_info.value.original_error, ValueError) |
||||
|
|
||||
|
@pytest.mark.anyio |
||||
|
async def test_mixed_patterns_warning(self): |
||||
|
"""Test warning when mixing lifespan and on_event.""" |
||||
|
|
||||
|
@asynccontextmanager |
||||
|
async def lifespan(app): |
||||
|
yield |
||||
|
|
||||
|
with pytest.warns(FastAPIDeprecationWarning, match="not recommended"): |
||||
|
UnifiedLifespanManager( |
||||
|
on_startup=[lambda: None], |
||||
|
lifespan=lifespan, |
||||
|
) |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# State Merging Tests |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class TestStateMerging: |
||||
|
"""Tests for lifespan state merging.""" |
||||
|
|
||||
|
def test_merge_child_overrides_parent(self): |
||||
|
"""Test default behavior: child overrides parent.""" |
||||
|
parent = {"shared": "parent", "parent_only": True} |
||||
|
child = {"shared": "child", "child_only": True} |
||||
|
|
||||
|
result = merge_lifespan_state(parent, child, child_overrides_parent=True) |
||||
|
|
||||
|
assert result["shared"] == "child" # child wins |
||||
|
assert result["parent_only"] is True |
||||
|
assert result["child_only"] is True |
||||
|
|
||||
|
def test_merge_parent_overrides_child(self): |
||||
|
"""Test alternative behavior: parent overrides child.""" |
||||
|
parent = {"shared": "parent", "parent_only": True} |
||||
|
child = {"shared": "child", "child_only": True} |
||||
|
|
||||
|
result = merge_lifespan_state(parent, child, child_overrides_parent=False) |
||||
|
|
||||
|
assert result["shared"] == "parent" # parent wins |
||||
|
assert result["parent_only"] is True |
||||
|
assert result["child_only"] is True |
||||
|
|
||||
|
def test_merge_none_states(self): |
||||
|
"""Test merging with None states.""" |
||||
|
assert merge_lifespan_state(None, None) == {} |
||||
|
assert merge_lifespan_state({"a": 1}, None) == {"a": 1} |
||||
|
assert merge_lifespan_state(None, {"b": 2}) == {"b": 2} |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# FastAPI Integration Tests |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class TestFastAPILifespanIntegration: |
||||
|
"""Integration tests with FastAPI app.""" |
||||
|
|
||||
|
def test_lifespan_context_manager(self): |
||||
|
"""Test modern lifespan context manager pattern.""" |
||||
|
events = [] |
||||
|
|
||||
|
@asynccontextmanager |
||||
|
async def lifespan(app: FastAPI) -> AsyncGenerator[dict[str, Any], None]: |
||||
|
events.append("startup") |
||||
|
yield {"db": "connection"} |
||||
|
events.append("shutdown") |
||||
|
|
||||
|
app = FastAPI(lifespan=lifespan) |
||||
|
|
||||
|
@app.get("/") |
||||
|
def index(): |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
with TestClient(app) as client: |
||||
|
events.append("running") |
||||
|
response = client.get("/") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
assert events == ["startup", "running", "shutdown"] |
||||
|
|
||||
|
def test_on_event_deprecation_warning(self): |
||||
|
"""Test that on_event emits deprecation warning.""" |
||||
|
app = FastAPI() |
||||
|
|
||||
|
with warnings.catch_warnings(record=True) as w: |
||||
|
warnings.simplefilter("always") |
||||
|
|
||||
|
@app.on_event("startup") |
||||
|
def startup(): |
||||
|
pass |
||||
|
|
||||
|
# Check that at least one warning was raised |
||||
|
assert len(w) >= 1 |
||||
|
# Check that one of them is about on_event deprecation |
||||
|
deprecation_warnings = [ |
||||
|
warning for warning in w |
||||
|
if "on_event" in str(warning.message) and "deprecated" in str(warning.message) |
||||
|
] |
||||
|
assert len(deprecation_warnings) >= 1 |
||||
|
|
||||
|
def test_on_event_handlers_still_work(self): |
||||
|
"""Test that on_event handlers still execute (backwards compatibility).""" |
||||
|
events = [] |
||||
|
app = FastAPI() |
||||
|
|
||||
|
# Suppress the deprecation warnings for this test |
||||
|
with warnings.catch_warnings(): |
||||
|
warnings.simplefilter("ignore") |
||||
|
|
||||
|
@app.on_event("startup") |
||||
|
def startup(): |
||||
|
events.append("startup") |
||||
|
|
||||
|
@app.on_event("shutdown") |
||||
|
def shutdown(): |
||||
|
events.append("shutdown") |
||||
|
|
||||
|
@app.get("/") |
||||
|
def index(): |
||||
|
return {"status": "ok"} |
||||
|
|
||||
|
with TestClient(app) as client: |
||||
|
events.append("running") |
||||
|
response = client.get("/") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
assert "startup" in events |
||||
|
assert "running" in events |
||||
|
assert "shutdown" in events |
||||
|
|
||||
|
def test_state_accessible_in_endpoints(self): |
||||
|
"""Test that lifespan state is accessible in endpoints.""" |
||||
|
|
||||
|
@asynccontextmanager |
||||
|
async def lifespan(app: FastAPI) -> AsyncGenerator[dict[str, Any], None]: |
||||
|
yield {"db_connection": "active"} |
||||
|
|
||||
|
app = FastAPI(lifespan=lifespan) |
||||
|
|
||||
|
@app.get("/state") |
||||
|
def get_state(request: Request): |
||||
|
# State should be accessible via request.state |
||||
|
return {"has_state": hasattr(request.state, "db_connection")} |
||||
|
|
||||
|
with TestClient(app) as client: |
||||
|
response = client.get("/state") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# create_unified_lifespan Tests |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class TestCreateUnifiedLifespan: |
||||
|
"""Tests for create_unified_lifespan function.""" |
||||
|
|
||||
|
@pytest.mark.anyio |
||||
|
async def test_create_from_lifespan_only(self): |
||||
|
"""Test creating unified lifespan from lifespan context only.""" |
||||
|
events = [] |
||||
|
|
||||
|
@asynccontextmanager |
||||
|
async def lifespan(app): |
||||
|
events.append("start") |
||||
|
yield {"key": "value"} |
||||
|
events.append("stop") |
||||
|
|
||||
|
unified = create_unified_lifespan(lifespan=lifespan) |
||||
|
|
||||
|
with warnings.catch_warnings(): |
||||
|
warnings.simplefilter("ignore", FastAPIDeprecationWarning) |
||||
|
async with unified(None) as state: |
||||
|
events.append("running") |
||||
|
assert state == {"key": "value"} |
||||
|
|
||||
|
assert events == ["start", "running", "stop"] |
||||
|
|
||||
|
@pytest.mark.anyio |
||||
|
async def test_create_from_handlers_only(self): |
||||
|
"""Test creating unified lifespan from handlers only.""" |
||||
|
events = [] |
||||
|
|
||||
|
async def on_start(): |
||||
|
events.append("start") |
||||
|
|
||||
|
async def on_stop(): |
||||
|
events.append("stop") |
||||
|
|
||||
|
unified = create_unified_lifespan( |
||||
|
on_startup=[on_start], |
||||
|
on_shutdown=[on_stop], |
||||
|
) |
||||
|
|
||||
|
with warnings.catch_warnings(): |
||||
|
warnings.simplefilter("ignore", FastAPIDeprecationWarning) |
||||
|
async with unified(None) as state: |
||||
|
events.append("running") |
||||
|
|
||||
|
assert "start" in events |
||||
|
assert "running" in events |
||||
|
assert "stop" in events |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# LifespanError Tests |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class TestLifespanError: |
||||
|
"""Tests for LifespanError exception.""" |
||||
|
|
||||
|
def test_error_message_formatting(self): |
||||
|
"""Test error message includes relevant information.""" |
||||
|
|
||||
|
def my_handler(): |
||||
|
pass |
||||
|
|
||||
|
error = LifespanError( |
||||
|
"Test error", |
||||
|
phase="startup", |
||||
|
handler=my_handler, |
||||
|
original_error=ValueError("Original"), |
||||
|
) |
||||
|
|
||||
|
error_str = str(error) |
||||
|
assert "startup" in error_str |
||||
|
assert "my_handler" in error_str |
||||
|
assert "Original" in error_str |
||||
|
|
||||
|
def test_error_attributes(self): |
||||
|
"""Test error attributes are accessible.""" |
||||
|
original = ValueError("test") |
||||
|
|
||||
|
error = LifespanError( |
||||
|
"Test", |
||||
|
phase="shutdown", |
||||
|
original_error=original, |
||||
|
) |
||||
|
|
||||
|
assert error.phase == "shutdown" |
||||
|
assert error.original_error is original |
||||
@ -0,0 +1,488 @@ |
|||||
|
""" |
||||
|
Tests for the OpenAPI plugin system. |
||||
|
""" |
||||
|
|
||||
|
from typing import Any |
||||
|
|
||||
|
import pytest |
||||
|
from fastapi import FastAPI |
||||
|
from fastapi.openapi.plugins import ( |
||||
|
OpenAPIPluginBase, |
||||
|
OpenAPIPluginSettings, |
||||
|
PluginExecutor, |
||||
|
PluginRegistry, |
||||
|
create_plugin, |
||||
|
) |
||||
|
from fastapi.testclient import TestClient |
||||
|
from pydantic import BaseModel |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# Test Models |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class Item(BaseModel): |
||||
|
name: str |
||||
|
price: float |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# Test Plugins |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class CustomExtensionPlugin(OpenAPIPluginBase): |
||||
|
"""Plugin that adds custom x- extensions to all operations.""" |
||||
|
|
||||
|
@property |
||||
|
def name(self) -> str: |
||||
|
return "custom-extension" |
||||
|
|
||||
|
@property |
||||
|
def priority(self) -> int: |
||||
|
return 50 # Higher priority (runs first) |
||||
|
|
||||
|
def modify_operation( |
||||
|
self, |
||||
|
route: Any, |
||||
|
method: str, |
||||
|
operation: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
operation["x-custom-extension"] = True |
||||
|
operation["x-route-name"] = route.name |
||||
|
return operation |
||||
|
|
||||
|
|
||||
|
class SchemaModifierPlugin(OpenAPIPluginBase): |
||||
|
"""Plugin that modifies the schema after generation.""" |
||||
|
|
||||
|
@property |
||||
|
def name(self) -> str: |
||||
|
return "schema-modifier" |
||||
|
|
||||
|
def modify_schema(self, schema: dict[str, Any]) -> dict[str, Any]: |
||||
|
schema["x-generated-by"] = "FastAPI Plugin System" |
||||
|
if "info" in schema: |
||||
|
schema["info"]["x-plugin-version"] = "1.0" |
||||
|
return schema |
||||
|
|
||||
|
|
||||
|
class PostGenerationPlugin(OpenAPIPluginBase): |
||||
|
"""Plugin that runs after schema generation.""" |
||||
|
|
||||
|
@property |
||||
|
def name(self) -> str: |
||||
|
return "post-generation" |
||||
|
|
||||
|
@property |
||||
|
def priority(self) -> int: |
||||
|
return 200 # Lower priority (runs last) |
||||
|
|
||||
|
def post_schema_generation(self, schema: dict[str, Any]) -> dict[str, Any]: |
||||
|
schema["x-post-processed"] = True |
||||
|
return schema |
||||
|
|
||||
|
|
||||
|
class PreSchemaPlugin(OpenAPIPluginBase): |
||||
|
"""Plugin that tracks pre_schema_generation calls.""" |
||||
|
|
||||
|
def __init__(self): |
||||
|
self.pre_called = False |
||||
|
self.settings_received: OpenAPIPluginSettings | None = None |
||||
|
|
||||
|
@property |
||||
|
def name(self) -> str: |
||||
|
return "pre-schema" |
||||
|
|
||||
|
def pre_schema_generation( |
||||
|
self, |
||||
|
app: FastAPI, |
||||
|
settings: OpenAPIPluginSettings, |
||||
|
) -> None: |
||||
|
self.pre_called = True |
||||
|
self.settings_received = settings |
||||
|
|
||||
|
|
||||
|
class ErrorPlugin(OpenAPIPluginBase): |
||||
|
"""Plugin that raises errors for testing error handling.""" |
||||
|
|
||||
|
@property |
||||
|
def name(self) -> str: |
||||
|
return "error-plugin" |
||||
|
|
||||
|
def modify_operation( |
||||
|
self, |
||||
|
route: Any, |
||||
|
method: str, |
||||
|
operation: dict[str, Any], |
||||
|
) -> dict[str, Any]: |
||||
|
raise ValueError("Intentional error for testing") |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# Plugin Registry Tests |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class TestPluginRegistry: |
||||
|
"""Tests for PluginRegistry class.""" |
||||
|
|
||||
|
def test_register_plugin(self): |
||||
|
"""Test basic plugin registration.""" |
||||
|
registry = PluginRegistry() |
||||
|
plugin = CustomExtensionPlugin() |
||||
|
registry.register(plugin) |
||||
|
|
||||
|
assert "custom-extension" in registry |
||||
|
assert len(registry) == 1 |
||||
|
assert registry.get_plugin("custom-extension") is plugin |
||||
|
|
||||
|
def test_register_duplicate_raises(self): |
||||
|
"""Test that registering duplicate name raises error.""" |
||||
|
registry = PluginRegistry() |
||||
|
plugin1 = CustomExtensionPlugin() |
||||
|
plugin2 = CustomExtensionPlugin() |
||||
|
|
||||
|
registry.register(plugin1) |
||||
|
with pytest.raises(ValueError, match="already registered"): |
||||
|
registry.register(plugin2) |
||||
|
|
||||
|
def test_register_with_replace(self): |
||||
|
"""Test replacing an existing plugin.""" |
||||
|
registry = PluginRegistry() |
||||
|
plugin1 = CustomExtensionPlugin() |
||||
|
plugin2 = CustomExtensionPlugin() |
||||
|
|
||||
|
registry.register(plugin1) |
||||
|
registry.register(plugin2, replace=True) |
||||
|
|
||||
|
assert registry.get_plugin("custom-extension") is plugin2 |
||||
|
|
||||
|
def test_unregister_plugin(self): |
||||
|
"""Test unregistering a plugin.""" |
||||
|
registry = PluginRegistry() |
||||
|
plugin = CustomExtensionPlugin() |
||||
|
registry.register(plugin) |
||||
|
|
||||
|
assert registry.unregister("custom-extension") is True |
||||
|
assert "custom-extension" not in registry |
||||
|
assert registry.unregister("custom-extension") is False |
||||
|
|
||||
|
def test_enable_disable_plugin(self): |
||||
|
"""Test enabling and disabling plugins.""" |
||||
|
registry = PluginRegistry() |
||||
|
plugin = CustomExtensionPlugin() |
||||
|
registry.register(plugin, enabled=False) |
||||
|
|
||||
|
assert not registry.is_enabled("custom-extension") |
||||
|
assert len(registry.get_active_plugins()) == 0 |
||||
|
|
||||
|
registry.enable("custom-extension") |
||||
|
assert registry.is_enabled("custom-extension") |
||||
|
assert len(registry.get_active_plugins()) == 1 |
||||
|
|
||||
|
registry.disable("custom-extension") |
||||
|
assert not registry.is_enabled("custom-extension") |
||||
|
|
||||
|
def test_priority_ordering(self): |
||||
|
"""Test that plugins are sorted by priority.""" |
||||
|
registry = PluginRegistry() |
||||
|
|
||||
|
low_priority = PostGenerationPlugin() # priority 200 |
||||
|
high_priority = CustomExtensionPlugin() # priority 50 |
||||
|
normal_priority = SchemaModifierPlugin() # priority 100 (default) |
||||
|
|
||||
|
registry.register(low_priority) |
||||
|
registry.register(high_priority) |
||||
|
registry.register(normal_priority) |
||||
|
|
||||
|
active = registry.get_active_plugins() |
||||
|
assert active[0].name == "custom-extension" # priority 50 |
||||
|
assert active[1].name == "schema-modifier" # priority 100 |
||||
|
assert active[2].name == "post-generation" # priority 200 |
||||
|
|
||||
|
def test_generation_count_increments(self): |
||||
|
"""Test that generation count increments on changes.""" |
||||
|
registry = PluginRegistry() |
||||
|
initial_count = registry.generation_count |
||||
|
|
||||
|
plugin = CustomExtensionPlugin() |
||||
|
registry.register(plugin) |
||||
|
assert registry.generation_count == initial_count + 1 |
||||
|
|
||||
|
registry.disable("custom-extension") |
||||
|
assert registry.generation_count == initial_count + 2 |
||||
|
|
||||
|
registry.enable("custom-extension") |
||||
|
assert registry.generation_count == initial_count + 3 |
||||
|
|
||||
|
registry.unregister("custom-extension") |
||||
|
assert registry.generation_count == initial_count + 4 |
||||
|
|
||||
|
def test_cache_invalidation_callback(self): |
||||
|
"""Test that cache invalidation callback is called.""" |
||||
|
callback_called = [] |
||||
|
|
||||
|
def on_invalidate(): |
||||
|
callback_called.append(True) |
||||
|
|
||||
|
registry = PluginRegistry(on_cache_invalidation=on_invalidate) |
||||
|
plugin = CustomExtensionPlugin() |
||||
|
|
||||
|
registry.register(plugin) |
||||
|
assert len(callback_called) == 1 |
||||
|
|
||||
|
registry.disable("custom-extension") |
||||
|
assert len(callback_called) == 2 |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# Plugin Execution Tests |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class TestPluginExecution: |
||||
|
"""Tests for plugin execution during schema generation.""" |
||||
|
|
||||
|
def test_modify_operation_plugin(self): |
||||
|
"""Test that modify_operation hook is called.""" |
||||
|
plugin = CustomExtensionPlugin() |
||||
|
app = FastAPI( |
||||
|
title="Test API", |
||||
|
version="1.0.0", |
||||
|
openapi_plugins=[plugin], |
||||
|
) |
||||
|
|
||||
|
@app.get("/items") |
||||
|
def get_items(): |
||||
|
return [] |
||||
|
|
||||
|
schema = app.openapi() |
||||
|
|
||||
|
# Check the custom extension was added |
||||
|
operation = schema["paths"]["/items"]["get"] |
||||
|
assert operation.get("x-custom-extension") is True |
||||
|
assert operation.get("x-route-name") == "get_items" |
||||
|
|
||||
|
def test_modify_schema_plugin(self): |
||||
|
"""Test that modify_schema hook is called.""" |
||||
|
plugin = SchemaModifierPlugin() |
||||
|
app = FastAPI( |
||||
|
title="Test API", |
||||
|
version="1.0.0", |
||||
|
openapi_plugins=[plugin], |
||||
|
) |
||||
|
|
||||
|
@app.get("/items") |
||||
|
def get_items(): |
||||
|
return [] |
||||
|
|
||||
|
schema = app.openapi() |
||||
|
|
||||
|
assert schema.get("x-generated-by") == "FastAPI Plugin System" |
||||
|
assert schema["info"].get("x-plugin-version") == "1.0" |
||||
|
|
||||
|
def test_post_generation_plugin(self): |
||||
|
"""Test that post_schema_generation hook is called.""" |
||||
|
plugin = PostGenerationPlugin() |
||||
|
app = FastAPI( |
||||
|
title="Test API", |
||||
|
version="1.0.0", |
||||
|
openapi_plugins=[plugin], |
||||
|
) |
||||
|
|
||||
|
@app.get("/items") |
||||
|
def get_items(): |
||||
|
return [] |
||||
|
|
||||
|
schema = app.openapi() |
||||
|
|
||||
|
assert schema.get("x-post-processed") is True |
||||
|
|
||||
|
def test_pre_schema_generation_plugin(self): |
||||
|
"""Test that pre_schema_generation hook is called.""" |
||||
|
plugin = PreSchemaPlugin() |
||||
|
app = FastAPI( |
||||
|
title="Test API", |
||||
|
version="1.0.0", |
||||
|
openapi_plugins=[plugin], |
||||
|
) |
||||
|
|
||||
|
@app.get("/items") |
||||
|
def get_items(): |
||||
|
return [] |
||||
|
|
||||
|
schema = app.openapi() |
||||
|
|
||||
|
assert plugin.pre_called is True |
||||
|
assert plugin.settings_received is not None |
||||
|
assert plugin.settings_received.title == "Test API" |
||||
|
assert plugin.settings_received.version == "1.0.0" |
||||
|
|
||||
|
def test_multiple_plugins(self): |
||||
|
"""Test multiple plugins working together.""" |
||||
|
app = FastAPI( |
||||
|
title="Test API", |
||||
|
version="1.0.0", |
||||
|
openapi_plugins=[ |
||||
|
CustomExtensionPlugin(), |
||||
|
SchemaModifierPlugin(), |
||||
|
PostGenerationPlugin(), |
||||
|
], |
||||
|
) |
||||
|
|
||||
|
@app.get("/items") |
||||
|
def get_items(): |
||||
|
return [] |
||||
|
|
||||
|
schema = app.openapi() |
||||
|
|
||||
|
# All plugins should have modified the schema |
||||
|
operation = schema["paths"]["/items"]["get"] |
||||
|
assert operation.get("x-custom-extension") is True |
||||
|
assert schema.get("x-generated-by") == "FastAPI Plugin System" |
||||
|
assert schema.get("x-post-processed") is True |
||||
|
|
||||
|
def test_plugin_error_handling(self): |
||||
|
"""Test that plugin errors don't break schema generation.""" |
||||
|
error_plugin = ErrorPlugin() |
||||
|
normal_plugin = SchemaModifierPlugin() |
||||
|
|
||||
|
app = FastAPI( |
||||
|
title="Test API", |
||||
|
version="1.0.0", |
||||
|
openapi_plugins=[error_plugin, normal_plugin], |
||||
|
) |
||||
|
|
||||
|
@app.get("/items") |
||||
|
def get_items(): |
||||
|
return [] |
||||
|
|
||||
|
# Should not raise, error is caught |
||||
|
with pytest.warns(UserWarning, match="raised an error"): |
||||
|
schema = app.openapi() |
||||
|
|
||||
|
# Normal plugin should still work |
||||
|
assert schema.get("x-generated-by") == "FastAPI Plugin System" |
||||
|
|
||||
|
def test_plugin_cache_invalidation(self): |
||||
|
"""Test that schema is regenerated when plugins change.""" |
||||
|
app = FastAPI( |
||||
|
title="Test API", |
||||
|
version="1.0.0", |
||||
|
) |
||||
|
|
||||
|
@app.get("/items") |
||||
|
def get_items(): |
||||
|
return [] |
||||
|
|
||||
|
# First generation - no plugins |
||||
|
schema1 = app.openapi() |
||||
|
assert schema1.get("x-generated-by") is None |
||||
|
|
||||
|
# Register a plugin |
||||
|
app.openapi_plugins.register(SchemaModifierPlugin()) |
||||
|
|
||||
|
# Should regenerate with plugin |
||||
|
schema2 = app.openapi() |
||||
|
assert schema2.get("x-generated-by") == "FastAPI Plugin System" |
||||
|
|
||||
|
# Disable plugin |
||||
|
app.openapi_plugins.disable("schema-modifier") |
||||
|
|
||||
|
# Should regenerate without plugin effect |
||||
|
schema3 = app.openapi() |
||||
|
assert schema3.get("x-generated-by") is None |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# create_plugin() Tests |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class TestCreatePlugin: |
||||
|
"""Tests for the create_plugin convenience function.""" |
||||
|
|
||||
|
def test_create_simple_plugin(self): |
||||
|
"""Test creating a plugin from callbacks.""" |
||||
|
plugin = create_plugin( |
||||
|
"simple-plugin", |
||||
|
modify_operation=lambda route, method, op: {**op, "x-simple": True}, |
||||
|
) |
||||
|
|
||||
|
assert plugin.name == "simple-plugin" |
||||
|
assert plugin.priority == 100 |
||||
|
|
||||
|
# Test the callback works |
||||
|
result = plugin.modify_operation(None, "GET", {"summary": "Test"}) |
||||
|
assert result["x-simple"] is True |
||||
|
assert result["summary"] == "Test" |
||||
|
|
||||
|
def test_create_plugin_with_priority(self): |
||||
|
"""Test creating a plugin with custom priority.""" |
||||
|
plugin = create_plugin( |
||||
|
"priority-plugin", |
||||
|
priority=25, |
||||
|
) |
||||
|
|
||||
|
assert plugin.priority == 25 |
||||
|
|
||||
|
|
||||
|
# ============================================================================= |
||||
|
# Integration Tests |
||||
|
# ============================================================================= |
||||
|
|
||||
|
|
||||
|
class TestIntegration: |
||||
|
"""Integration tests with actual HTTP requests.""" |
||||
|
|
||||
|
def test_plugin_with_test_client(self): |
||||
|
"""Test plugin output via test client.""" |
||||
|
app = FastAPI( |
||||
|
title="Test API", |
||||
|
version="1.0.0", |
||||
|
openapi_plugins=[CustomExtensionPlugin()], |
||||
|
) |
||||
|
|
||||
|
@app.get("/items", response_model=list[Item]) |
||||
|
def get_items(): |
||||
|
return [] |
||||
|
|
||||
|
@app.post("/items", response_model=Item) |
||||
|
def create_item(item: Item): |
||||
|
return item |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
response = client.get("/openapi.json") |
||||
|
assert response.status_code == 200 |
||||
|
|
||||
|
schema = response.json() |
||||
|
|
||||
|
# Check both operations have custom extension |
||||
|
assert schema["paths"]["/items"]["get"].get("x-custom-extension") is True |
||||
|
assert schema["paths"]["/items"]["post"].get("x-custom-extension") is True |
||||
|
|
||||
|
def test_docs_with_plugins(self): |
||||
|
"""Test that Swagger UI still works with plugins.""" |
||||
|
app = FastAPI( |
||||
|
title="Test API", |
||||
|
version="1.0.0", |
||||
|
openapi_plugins=[SchemaModifierPlugin()], |
||||
|
) |
||||
|
|
||||
|
@app.get("/items") |
||||
|
def get_items(): |
||||
|
return [] |
||||
|
|
||||
|
client = TestClient(app) |
||||
|
|
||||
|
# Swagger UI should load |
||||
|
response = client.get("/docs") |
||||
|
assert response.status_code == 200 |
||||
|
assert "swagger-ui" in response.text.lower() |
||||
|
|
||||
|
# ReDoc should load |
||||
|
response = client.get("/redoc") |
||||
|
assert response.status_code == 200 |
||||
|
assert "redoc" in response.text.lower() |
||||
Loading…
Reference in new issue