diff --git a/tofupilot/client.py b/tofupilot/client.py index b6f91d6..56e7427 100644 --- a/tofupilot/client.py +++ b/tofupilot/client.py @@ -7,6 +7,8 @@ from datetime import datetime, timedelta from importlib.metadata import version +import json +import base64 import requests from .constants import ( @@ -27,6 +29,7 @@ handle_response, handle_http_error, handle_network_error, + notify_server, ) @@ -65,6 +68,7 @@ def create_run( # pylint: disable=too-many-arguments,too-many-locals run_passed: bool, procedure_id: Optional[str] = None, procedure_name: Optional[str] = None, + procedure_version: Optional[str] = None, steps: Optional[List[Step]] = None, phases: Optional[List[Phase]] = None, started_at: Optional[datetime] = None, @@ -85,6 +89,8 @@ def create_run( # pylint: disable=too-many-arguments,too-many-locals The unique identifier of the procedure to which the test run belongs. Required if several procedures exists with the same procedure_name. procedure_name (str, optional): The name of the procedure to which the test run belongs. A new procedure will be created if none was found with this name. + procedure_version (str, optional): + The version of the procedure to which the test run belongs. started_at (datetime, optional): The datetime at which the test started. Default is None. duration (timedelta, optional): @@ -117,6 +123,7 @@ def create_run( # pylint: disable=too-many-arguments,too-many-locals "run_passed": run_passed, "procedure_id": procedure_id, "procedure_name": procedure_name, + "procedure_version": procedure_version, "client": "Python", "client_version": self._current_version, } @@ -169,10 +176,7 @@ def create_run( # pylint: disable=too-many-arguments,too-many-locals except requests.RequestException as e: return handle_network_error(self._logger, e) - def create_run_from_openhtf_report( - self, - file_path: str, - ) -> dict: + def create_run_from_openhtf_report(self, file_path: str): """ Creates a run on TofuPilot from an OpenHTF JSON report. @@ -187,49 +191,70 @@ def create_run_from_openhtf_report( References: https://www.tofupilot.com/docs/api#create-a-run-from-a-file """ - self._logger.info("Starting run creation...") - - # Validate report - validate_files( - self._logger, [file_path], self._max_attachments, self._max_file_size - ) - - # Upload report - try: - upload_id = upload_file(self._headers, self._url, file_path) - except requests.exceptions.HTTPError as http_err: - return handle_http_error(self._logger, http_err) - except requests.RequestException as e: - return handle_network_error(self._logger, e) - - payload = { - "upload_id": upload_id, - "importer": "OPENHTF", - "client": "Python", - "client_version": self._current_version, - } + # Upload report and create run from file_path + run_id = self.upload_and_create_from_openhtf_report(file_path) - self._log_request("POST", "/import", payload) - - # Create run from file try: - response = requests.post( - f"{self._url}/import", - json=payload, - headers=self._headers, - timeout=SECONDS_BEFORE_TIMEOUT, - ) - response.raise_for_status() - result = handle_response(self._logger, response) - - run_id = result.get("id") - - return run_id - - except requests.exceptions.HTTPError as http_err: - return handle_http_error(self._logger, http_err) - except requests.RequestException as e: - return handle_network_error(self._logger, e) + with open(file_path, "r", encoding="utf-8") as file: + test_record = json.load(file) + except FileNotFoundError: + print(f"Error: The file '{file_path}' was not found.") + except json.JSONDecodeError: + print(f"Error: The file '{file_path}' contains invalid JSON.") + except PermissionError: + print(f"Error: Insufficient permissions to read '{file_path}'.") + except Exception as e: + print(f"Unexpected error: {e}") + + if run_id and test_record: + number_of_attachments = 0 + for phase in test_record.get("phases"): + # Keep only max number of attachments + if number_of_attachments >= self._max_attachments: + self._logger.warning( + "Too many attachments, trimming to %d attachments.", + self._max_attachments, + ) + break + for attachment_name, attachment in phase.get("attachments").items(): + number_of_attachments += 1 + + self._logger.info("Uploading %s...", attachment_name) + + # Upload initialization + initialize_url = f"{self._url}/uploads/initialize" + payload = {"name": attachment_name} + + response = requests.post( + initialize_url, + data=json.dumps(payload), + headers=self._headers, + timeout=SECONDS_BEFORE_TIMEOUT, + ) + + response.raise_for_status() + response_json = response.json() + upload_url = response_json.get("uploadUrl") + upload_id = response_json.get("id") + + data = base64.b64decode(attachment["data"]) + + requests.put( + upload_url, + data=data, + headers={ + "Content-Type": attachment["mimetype"] + or "application/octet-stream", # Default to binary if mimetype is missing + }, + timeout=SECONDS_BEFORE_TIMEOUT, + ) + + notify_server(self._headers, self._url, upload_id, run_id) + + self._logger.success( + "Attachment %s successfully uploaded and linked to run.", + attachment_name, + ) def get_runs(self, serial_number: str) -> dict: """ @@ -384,6 +409,62 @@ def delete_unit(self, serial_number: str) -> dict: except requests.RequestException as e: return handle_network_error(self._logger, e) + def upload_and_create_from_openhtf_report( + self, + file_path: str, + ) -> str: + """ + Takes a path to an OpenHTF JSON file report, uploads it and creates a run from it. + + Returns: + str: + Id of the newly created run + """ + + self._logger.info("Starting run creation...") + + # Validate report + validate_files( + self._logger, [file_path], self._max_attachments, self._max_file_size + ) + + # Upload report + try: + upload_id = upload_file(self._headers, self._url, file_path) + except requests.exceptions.HTTPError as http_err: + return handle_http_error(self._logger, http_err) + except requests.RequestException as e: + return handle_network_error(self._logger, e) + + payload = { + "upload_id": upload_id, + "importer": "OPENHTF", + "client": "Python", + "client_version": self._current_version, + } + + self._log_request("POST", "/import", payload) + + # Create run from file + try: + response = requests.post( + f"{self._url}/import", + json=payload, + headers=self._headers, + timeout=SECONDS_BEFORE_TIMEOUT, + ) + response.raise_for_status() + result = handle_response(self._logger, response) + + run_id = result.get("id") + + return run_id + + except requests.exceptions.HTTPError as http_err: + return handle_http_error(self._logger, http_err) + except requests.RequestException as e: + return handle_network_error(self._logger, e) + def get_websocket_url(self) -> dict: """ Fetches websocket connection url associated with API Key. diff --git a/tofupilot/openhtf/upload.py b/tofupilot/openhtf/upload.py index f225c4f..42a3a93 100644 --- a/tofupilot/openhtf/upload.py +++ b/tofupilot/openhtf/upload.py @@ -22,13 +22,13 @@ class upload: # pylint: disable=invalid-name OpenHTF output callback to automatically upload the test report to TofuPilot upon test completion. This function behaves similarly to manually parsing the OpenHTF JSON test report and calling - `TofuPilotClient().create_run()` with the parsed data, streamlining the process for automatic uploads. + `TofuPilotClient().create_run()` with the parsed data. ### Usage Example: ```python from openhtf import test - import tofupilot + from tofupilot.openhtf import upload # ... @@ -36,7 +36,7 @@ def main(): test = Test(*your_phases, procedure_id="FVT1") # Add TofuPilot's upload callback to automatically send the test report upon completion - test.add_output_callback(tofupilot.upload()) + test.add_output_callback(upload()) test.execute(lambda: "SN15") ``` @@ -69,9 +69,7 @@ def __call__(self, test_record: TestRecord): ) # Format the timestamp as YYYY-MM-DD_HH_MM_SS_SSS - start_time_formatted = start_time.strftime("%Y-%m-%d_%H-%M-%S-%f")[ - :-3 - ] # Use underscores for time, slice for milliseconds precision + start_time_formatted = start_time.strftime("%Y-%m-%d_%H-%M-%S-%f")[:-3] temp_dir = tempfile.gettempdir() @@ -95,7 +93,7 @@ def __call__(self, test_record: TestRecord): try: # Call create_run_from_report with the generated file path - run_id = self.client.create_run_from_openhtf_report(filename) + run_id = self.client.upload_and_create_from_openhtf_report(filename) finally: # Ensure the file is deleted after processing if os.path.exists(filename):