Coverage for src/paperap/plugins/collect_test_data.py: 93%
137 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-03-22 16:02 -0400
« prev ^ index » next coverage.py v7.6.12, created at 2025-03-22 16:02 -0400
1"""
3----------------------------------------------------------------------------
5 METADATA:
7 File: collect_test_data.py
8 Project: paperap
9 Created: 2025-03-04
10 Version: 0.0.10
11 Author: Jess Mann
12 Email: jess@jmann.me
13 Copyright (c) 2025 Jess Mann
15----------------------------------------------------------------------------
17 LAST MODIFIED:
19 2025-03-04 By Jess Mann
21"""
23from __future__ import annotations
25import datetime
26import json
27import logging
28import re
29from decimal import Decimal
30from pathlib import Path
31from typing import TYPE_CHECKING, Any, override
33from faker import Faker
34from pydantic import HttpUrl, field_validator
36from paperap.exceptions import ModelValidationError
37from paperap.models import StandardModel
38from paperap.plugins.base import Plugin
39from paperap.signals import SignalPriority, registry
41logger = logging.getLogger(__name__)
43sanitize_pattern = re.compile(r"[^a-zA-Z0-9|.=_-]")
45SANITIZE_KEYS = [
46 "email",
47 "first_name",
48 "last_name",
49 "name",
50 "phone",
51 "username",
52 "content",
53 "filename",
54 "title",
55 "slug",
56 "original_filename",
57 "archived_file_name",
58 "task_file_name",
59 "filename",
60]
62type ClientResponse = dict[str, Any] | list[dict[str, Any]]
65class SampleDataCollector(Plugin):
66 """
67 Plugin to collect test data from API responses.
68 """
70 name = "test_data_collector"
71 description = "Collects sample data from API responses for testing purposes"
72 version = "0.0.3"
73 fake: Faker = Faker()
74 test_dir: Path = Path("tests/sample_data")
76 @field_validator("test_dir", mode="before")
77 @classmethod
78 def validate_test_dir(cls, value: Any) -> Path | None:
79 """Validate the test directory path."""
80 # Convert string path to Path object if needed
81 if not value:
82 value = Path("tests/sample_data")
84 if isinstance(value, str):
85 value = Path(value)
87 if not isinstance(value, Path):
88 raise ModelValidationError("Test directory must be a string or Path object")
90 if not value.is_absolute():
91 # Make it relative to project root
92 project_root = Path(__file__).parents[4]
93 value = project_root / value
95 value.mkdir(parents=True, exist_ok=True)
96 return value
98 @override
99 def setup(self) -> None:
100 """Register signal handlers."""
101 registry.connect("resource._handle_response:after", self.save_list_response, SignalPriority.LOW)
102 registry.connect("resource._handle_results:before", self.save_first_item, SignalPriority.LOW)
103 registry.connect("client.request:after", self.save_parsed_response, SignalPriority.LOW)
105 @override
106 def teardown(self) -> None:
107 """Unregister signal handlers."""
108 registry.disconnect("resource._handle_response:after", self.save_list_response)
109 registry.disconnect("resource._handle_results:before", self.save_first_item)
110 registry.disconnect("client.request:after", self.save_parsed_response)
112 @staticmethod
113 def _json_serializer(obj: Any) -> Any:
114 """Serialize objects that are not natively serializable."""
115 if isinstance(obj, datetime.datetime):
116 return obj.isoformat()
117 if isinstance(obj, Path):
118 return str(obj)
119 if isinstance(obj, Decimal):
120 return float(obj)
121 if isinstance(obj, StandardModel):
122 return obj.to_dict()
123 if isinstance(obj, StandardModel):
124 return obj.model_dump()
125 if isinstance(obj, set):
126 return list(obj)
127 if isinstance(obj, bytes):
128 return obj.decode("utf-8")
129 raise TypeError(f"Type {type(obj).__name__} is not JSON serializable")
131 def _sanitize_list_response[R: list[dict[str, Any]]](self, response: R) -> R:
132 """
133 Sanitize the response data to replace any strings with potentially personal information with dummy data
134 """
135 sanitized_list: R = [] # type: ignore
136 for item in response:
137 sanitized_item = self._sanitize_value_recursive("", item)
138 sanitized_list.append(sanitized_item) # type: ignore
139 return sanitized_list
141 def _sanitize_dict_response[R: dict[str, Any]](self, **response: R) -> R:
142 """
143 Sanitize the response data to replace any strings with potentially personal information with dummy data
144 """
145 sanitized: dict[str, Any] = {}
146 for key, value in response.items():
147 sanitized[key] = self._sanitize_value_recursive(key, value)
149 # Replace "next" domain using regex
150 if (next_page := response.get("next", None)) and isinstance(next_page, str):
151 sanitized["next"] = re.sub(r"https?://.*?/", "https://example.com/", next_page)
153 return sanitized # type: ignore
155 def _sanitize_value_recursive(self, key: str, value: Any) -> Any:
156 """
157 Recursively sanitize the value to replace any strings with potentially personal information with dummy data
158 """
159 if isinstance(value, dict):
160 return {k: self._sanitize_value_recursive(k, v) for k, v in value.items()}
162 if key in SANITIZE_KEYS:
163 if isinstance(value, str):
164 return self.fake.word()
165 if isinstance(value, list):
166 return [self.fake.word() for _ in value]
168 return value
170 def save_response(self, filepath: Path, response: ClientResponse | None, **kwargs: Any) -> None:
171 """
172 Save the response to a JSON file.
173 """
174 if not response or filepath.exists():
175 return
177 try:
178 if isinstance(response, list):
179 response = self._sanitize_list_response(response)
180 else:
181 response = self._sanitize_dict_response(**response)
182 filepath.parent.mkdir(parents=True, exist_ok=True)
183 with filepath.open("w") as f:
184 json.dump(response, f, indent=4, sort_keys=True, ensure_ascii=False, default=self._json_serializer)
185 except (TypeError, OverflowError, OSError) as e:
186 # Don't allow the plugin to interfere with normal operations in the event of failure
187 logger.error("Error saving response to file (%s): %s", filepath.absolute(), e)
189 def save_list_response[R: ClientResponse | None](self, sender: Any, response: R, **kwargs: Any) -> R:
190 """Save the list response to a JSON file."""
191 if not response or not (resource_name := kwargs.get("resource")):
192 return response
194 filepath = self.test_dir / f"{resource_name}_list.json"
195 self.save_response(filepath, response)
197 return response
199 def save_first_item[R: dict[str, Any]](self, sender: Any, item: R, **kwargs: Any) -> R:
200 """Save the first item from a list to a JSON file."""
201 resource_name = kwargs.get("resource")
202 if not resource_name:
203 return item
205 filepath = self.test_dir / f"{resource_name}_item.json"
206 self.save_response(filepath, item)
208 # Disable this handler after saving the first item
209 registry.disable("resource._handle_results:before", self.save_first_item)
211 return item
213 def save_parsed_response(
214 self,
215 parsed_response: dict[str, Any],
216 method: str,
217 params: dict[str, Any] | None,
218 json_response: bool,
219 endpoint: str | HttpUrl,
220 **kwargs: Any,
221 ) -> dict[str, Any]:
222 """
223 Save the request data to a JSON file.
225 Connects to client.request:after signal.
226 """
227 if not endpoint:
228 raise ValueError("Endpoint is required to save parsed response")
230 endpoint = str(endpoint)
232 # If endpoint contains "example.com", we're testing, so skip it
233 if "example.com" in str(endpoint):
234 return parsed_response
236 if not json_response or not params:
237 return parsed_response
239 # Strip url to final path segment
240 resource_name = ".".join(endpoint.split("/")[-2:])
242 combined_params = list(f"{k}={v}" for k, v in params.items())
243 params_str = "|".join(combined_params)
244 filename_prefix = ""
245 if method.lower() != "get":
246 filename_prefix = f"{method.lower()}__"
247 filename = f"{filename_prefix}{resource_name}__{params_str}.json"
248 filename = sanitize_pattern.sub("_", filename)
249 filename = filename[:100] # Limit filename length
251 filepath = self.test_dir / filename
252 self.save_response(filepath, parsed_response)
254 return parsed_response
256 @override
257 @classmethod
258 def get_config_schema(cls) -> dict[str, Any]:
259 """Define the configuration schema for this plugin."""
260 return {
261 "test_dir": {
262 "type": str,
263 "description": "Directory to save test data files",
264 "required": False,
265 }
266 }