Coverage for graphqler / utils / objects_bucket.py: 68%
175 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-20 10:09 -0400
1"""Class for an objects bucket to contain the history of all objects in the system under test
2What does an objects bucket track?
3- For each type of object, track any values that were associated to that object to be used later
4- For each kind of scalar, track any values seen to be used later
6TODO: Implement the following:
7The class should have two functionalities
81. Given the graphql data response, parse the data and put objects in the bucket
92. Be able to return random scalars / objects from the bucket
103. Be able to return objects from the bucket if given a type and the object name
11"""
13import pathlib
14import pprint
15import random
16from typing import Self
18import cloudpickle as pickle
20from graphqler import config
21from graphqler.utils.api import API
22from graphqler.utils.file_utils import get_or_create_file
23from graphqler.utils.parser_utils import get_output_type_from_details
25from .singleton import singleton
28@singleton
29class ObjectsBucket:
30 def __init__(self, api: API):
31 self.api = api
33 # Stores {object_name: {type: str, results: dict}} where list is a result with the scalar fields of the object
34 self.objects: dict[str, list] = {}
36 # Stores the raw scalars {scalar_name: {type: str, values: set() }} where set() is a result with the scalar fields of the object
37 self.scalars: dict[str, dict] = {}
39 # File paths
40 self.pickle_save_path = pathlib.Path(config.OUTPUT_DIRECTORY) / config.SERIALIZED_DIR_NAME / config.OBJECTS_BUCKET_PICKLE_FILE_NAME
41 self.text_save_path = pathlib.Path(config.OUTPUT_DIRECTORY) / config.OBJECTS_BUCKET_TEXT_FILE_NAME
43 def __str__(self):
44 """Returns a string representation of the objects bucket"""
45 built_str = "\n------------------- OBJECTS BUCKET -------------------\n"
46 built_str += pprint.pformat(self.objects)
48 built_str += "\n\n"
49 built_str += "\n------------------- SCALARS BUCKET -------------------\n"
50 built_str += pprint.pformat(self.scalars)
52 return built_str
54 # ------------------- Pickle -------------------
55 def __getstate__(self):
56 # Return a dictionary of the attributes to pickle
57 return self.__dict__
59 def __setstate__(self, state):
60 # Restore the state from the pickled attributes
61 self.__dict__.update(state)
63 def save(self):
64 """Saves the objects bucket as a pickle file and as a text file"""
65 self.pickle_save_path = get_or_create_file(self.pickle_save_path)
66 with open(self.pickle_save_path, "wb") as file:
67 pickle.dump(self, file)
69 self.text_save_path = get_or_create_file(self.text_save_path)
70 with open(self.text_save_path, "w") as file:
71 file.write(f"Number of objects: {self.get_num_objects()}\n")
72 file.write(f"Number of scalars: {self.get_num_scalars()}\n")
73 file.write(str(self))
75 def load(self) -> Self:
76 """Loads the objects bucket from a pickle file. If the file doesn't exist, does nothing.
77 """
78 if self.pickle_save_path.exists():
79 with open(self.pickle_save_path, "rb") as file:
80 loaded_bucket = pickle.load(file)
81 self.__dict__ = loaded_bucket.__dict__
83 return self
85 # ------------------- GETTERS -------------------
86 def get_num_objects(self) -> int:
87 """Returns the number of objects in the bucket
89 Returns:
90 int: The number of objects in the bucket
91 """
92 sum = 0
93 for object_name, object_info in self.objects.items():
94 sum += len(object_info)
95 return sum
97 def get_num_scalars(self) -> int:
98 """Returns the number of scalars in the bucket
100 Returns:
101 int: The number of scalars in the bucket
102 """
103 sum = 0
104 for scalar_name, scalar_info in self.scalars.items():
105 sum += len(scalar_info["values"])
106 return sum
108 def get_random_object(self, object_name: str) -> dict:
109 """Returns a random object from the bucket
111 Args:
112 object_name (str): The object name
114 Returns:
115 dict: A random object from the bucket
116 """
117 if object_name not in self.objects:
118 return {}
120 return random.choice(self.objects[object_name])
122 def get_random_object_field_value(self, object_name: str, field_name: str) -> str | int | float | bool | None:
123 """Returns a random field from an object.
124 Retries up to 5 times to find a field value that isn't None
126 Args:
127 object_name (str): The object name
128 field_name (str): The field name
130 Returns:
131 str | int | float | bool: The field value
132 """
133 if object_name not in self.objects:
134 raise Exception("Object not found in bucket")
136 max_retries = 5
137 num_retries = 0
138 used_indices = []
139 while num_retries < max_retries:
140 length_of_objects = len(self.objects[object_name])
141 if len(used_indices) == length_of_objects:
142 return None
143 random_index = random.choice([i for i in range(length_of_objects) if i not in used_indices])
144 object_to_use = self.objects[object_name][random_index]
145 found_key, found_value = self.find_key_in_dict(object_to_use, field_name)
147 if found_value is not None:
148 return found_value
149 else:
150 num_retries += 1
151 used_indices.append(random_index)
152 return None
154 # ------------------- SETTERS -------------------
155 def put_in_bucket(self, response_data: dict) -> bool:
156 """Puts an object in the bucket, returns True if the object was added, False otherwise
158 Args:
159 response_data (dict): The data to put in the bucket. This is the responses data from GraphQL
161 Returns:
162 bool: True if the object was added, False otherwise
163 """
164 # If no data, just return
165 if not response_data:
166 return False
168 # Iterate through the data, put in the bucket
169 for data_key, data in response_data.items():
170 if self.api.is_operation_in_api(data_key):
171 self.parse_as_object(data_key, data)
172 # Regardless, always parse the entire data into our scalars bucket as well for future lookups
173 self.parse_as_scalar(data_key, data)
174 return True
176 def parse_as_object(self, operation_name: str, data: dict | list[dict]):
177 """Parses the data as an object by looking up the output of the operation in the API
179 Args:
180 operation_name (str): The operation name, should be an operation in the API
181 data (dict | List[dict]): The data to parse
182 """
183 # Get the operation from the API
184 operation = self.api.get_operation(operation_name)
185 operation_output_type = get_output_type_from_details(operation)
187 if isinstance(data, list):
188 for item in data:
189 self.put_object_in_bucket(operation_output_type, item)
190 else:
191 self.put_object_in_bucket(operation_output_type, data)
193 def parse_as_scalar(self, method_name: str, method_data: dict | list | str | int | float | bool | None):
194 """Parses the data as a scalar (can be a list, dict, or any of the base GraphQL types
196 Args:
197 method_name (str): The method name
198 method_data (str): The method data
199 """
200 if isinstance(method_data, str):
201 self.put_scalar_in_bucket(method_name, "String", method_data)
202 elif isinstance(method_data, bool):
203 self.put_scalar_in_bucket(method_name, "Boolean", method_data)
204 elif isinstance(method_data, int):
205 self.put_scalar_in_bucket(method_name, "Int", method_data)
206 elif isinstance(method_data, float):
207 self.put_scalar_in_bucket(method_name, "Float", method_data)
208 elif isinstance(method_data, list):
209 for item in method_data:
210 self.parse_as_scalar(method_name, item)
211 elif isinstance(method_data, dict):
212 self.parse_object_scalars(method_data)
214 def put_object_in_bucket(self, object_name: str, object_info: dict):
215 """Puts an object in the bucket, skipping duplicates
217 Args:
218 object_name (str): The object's name
219 object_info (dict): The object's info
220 """
221 if object_name not in self.objects:
222 self.objects[object_name] = []
223 if object_info not in self.objects[object_name]:
224 self.objects[object_name].append(object_info)
226 def parse_object_scalars(self, object_info: dict):
227 """Parses each field of a dictionary as a scalar and parses it into the scalar components
229 Args:
230 object_info (dict): The object info
231 """
232 for field_name, field_value in object_info.items():
233 self.parse_as_scalar(field_name, field_value)
235 def put_scalar_in_bucket(self, name: str, type: str, data: str | int | float | bool):
236 """Puts scalar in the bucket
238 Args:
239 name (str): The scalar's name
240 type (str): The scalar's type
241 data (str): The scalar's data
242 """
243 if name not in self.scalars:
244 self.scalars[name] = {"type": type, "values": {data}}
245 self.scalars[name]["values"].add(data)
247 # ------------------- DELETERS -------------------
248 def delete_object_from_bucket(self, object_name: str, object_value: dict):
249 """Deletes a specific object entry from the bucket
251 Args:
252 object_name (str): The object type name
253 object_value (dict): The specific object instance to remove
254 """
255 if object_name not in self.objects:
256 return
257 try:
258 self.objects[object_name].remove(object_value)
259 except ValueError:
260 pass
262 def update_object_in_bucket(self, response_data: dict):
263 """Updates an existing object in the bucket with new data from a successful UPDATE mutation response.
264 If the object is not already tracked, adds it.
266 Args:
267 response_data (dict): The data dict from the GraphQL response
268 """
269 if not response_data:
270 return
271 for data_key, data in response_data.items():
272 if not self.api.is_operation_in_api(data_key):
273 continue
274 operation = self.api.get_operation(data_key)
275 operation_output_type = get_output_type_from_details(operation)
276 if not isinstance(data, dict) or operation_output_type not in self.objects:
277 self.put_in_bucket(response_data)
278 return
279 # Try to find and replace the matching entry by id-like fields
280 id_value = data.get("id") or data.get("ID")
281 if id_value is not None:
282 for i, existing in enumerate(self.objects[operation_output_type]):
283 if existing.get("id") == id_value or existing.get("ID") == id_value:
284 self.objects[operation_output_type][i] = data
285 return
286 # No match found — just add it
287 self.put_object_in_bucket(operation_output_type, data)
289 # ------------------- HELPERS -------------------
290 def clear_bucket(self):
291 """Clears the bucket"""
292 self.objects.clear()
293 self.scalars.clear()
295 def is_empty(self) -> bool:
296 """Checks if the object bucket is empty
298 Returns:
299 bool: True if the object bucket is empty, False otherwise
300 """
301 return len(self.objects) == 0 and len(self.scalars) == 0
303 def is_object_in_bucket(self, object_name: str) -> bool:
304 """Checks if an object is in the bucket
306 Args:
307 object_name (str): The object name
309 Returns:
310 bool: True if the object is in the bucket, False otherwise
311 """
312 if not config.USE_OBJECTS_BUCKET:
313 return False
314 return object_name in self.objects and len(self.objects[object_name]) > 0
316 def find_key_in_dict(self, dictionary: dict, key: str) -> tuple[str, str | int | float | bool | None]:
317 """Recursively searches for the key in a nested dictionary and returns its full path and value.
319 Args:
320 dictionary (dict): The dictionary to search
321 key (str): The key to search for
323 Returns:
324 tuple[str, str | int | float | bool | None]: The key and value
326 """
327 if not config.USE_OBJECTS_BUCKET:
328 return ("", None)
329 for k, v in dictionary.items():
330 if k == key:
331 return k, v
332 if isinstance(v, dict):
333 result = self.find_key_in_dict(v, key)
334 if result is not None:
335 return result
336 return ("", None)
338 def get_random_scalar_from_bucket_by_type(self, scalar_type: str) -> str | int | float | bool:
339 """Gets a random scalar from the bucket
341 Args:
342 scalar_type (str): The scalar type
344 Returns:
345 str | int | float | bool: The scalar value
346 """
347 if not config.USE_OBJECTS_BUCKET:
348 return ""
349 for scalar_name, scalar in self.scalars.items():
350 if scalar["type"] == scalar_type:
351 return random.choice(list(scalar["values"]))
352 return ""
354 def get_random_scalar_from_bucket_by_name(self, scalar_name) -> str | int | float | bool:
355 """Gets a random scalar from the bucket with the name
357 Args:
358 scalar_name (str): The scalar name
360 Returns:
361 str | int | float | bool: The scalar value
362 """
363 if not config.USE_OBJECTS_BUCKET:
364 return ""
365 if scalar_name not in self.scalars:
366 return ""
368 return random.choice(list(self.scalars[scalar_name]["values"]))