diff --git a/tofupilot/client.py b/tofupilot/client.py index 56e7427..088880f 100644 --- a/tofupilot/client.py +++ b/tofupilot/client.py @@ -34,9 +34,23 @@ class TofuPilotClient: - """Wrapper for TofuPilot's API that provides additional support for handling attachments.""" + """Wrapper for TofuPilot's API that provides additional support for handling attachments. + + Args: + api_key (Optional[str]): API key for authentication with TofuPilot's API. + If not provided, the TOFUPILOT_API_KEY environment variable will be used. + url (Optional[str]): Base URL for TofuPilot's API. + If not provided, the TOFUPILOT_URL environment variable or the default endpoint will be used. + verify (Optional[str]): Path to a CA bundle file to verify TofuPilot's server certificate. + Useful for connecting to instances with custom/self-signed certificates. + """ - def __init__(self, api_key: Optional[str] = None, url: Optional[str] = None): + def __init__( + self, + api_key: Optional[str] = None, + url: Optional[str] = None, + verify: Optional[str] = None, + ): self._current_version = version("tofupilot") print_version_banner(self._current_version) self._logger = setup_logger(logging.INFO) @@ -52,6 +66,7 @@ def __init__(self, api_key: Optional[str] = None, url: Optional[str] = None): "Content-Type": "application/json", "Authorization": f"Bearer {self._api_key}", } + self._verify = verify self._max_attachments = CLIENT_MAX_ATTACHMENTS self._max_file_size = FILE_MAX_SIZE check_latest_version(self._logger, self._current_version, "tofupilot") @@ -159,6 +174,7 @@ def create_run( # pylint: disable=too-many-arguments,too-many-locals json=payload, headers=self._headers, timeout=SECONDS_BEFORE_TIMEOUT, + verify=self._verify, ) response.raise_for_status() result = handle_response(self._logger, response) @@ -166,7 +182,12 @@ def create_run( # pylint: disable=too-many-arguments,too-many-locals run_id = result.get("id") if run_id and attachments: upload_attachments( - self._logger, self._headers, self._url, attachments, run_id + self._logger, + self._headers, + self._url, + attachments, + run_id, + self._verify, ) return result @@ -229,6 +250,7 @@ def create_run_from_openhtf_report(self, file_path: str): initialize_url, data=json.dumps(payload), headers=self._headers, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -249,7 +271,13 @@ def create_run_from_openhtf_report(self, file_path: str): timeout=SECONDS_BEFORE_TIMEOUT, ) - notify_server(self._headers, self._url, upload_id, run_id) + notify_server( + self._headers, + self._url, + upload_id, + run_id, + self._verify, + ) self._logger.success( "Attachment %s successfully uploaded and linked to run.", @@ -293,6 +321,7 @@ def get_runs(self, serial_number: str) -> dict: response = requests.get( f"{self._url}/runs", headers=self._headers, + verify=self._verify, params=params, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -326,6 +355,7 @@ def delete_run(self, run_id: str) -> dict: response = requests.delete( f"{self._url}/runs/{run_id}", headers=self._headers, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -366,6 +396,7 @@ def update_unit( f"{self._url}/units/{serial_number}", json=payload, headers=self._headers, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -399,6 +430,7 @@ def delete_unit(self, serial_number: str) -> dict: response = requests.delete( f"{self._url}/units/{serial_number}", headers=self._headers, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -430,7 +462,9 @@ def upload_and_create_from_openhtf_report( # Upload report try: - upload_id = upload_file(self._headers, self._url, file_path) + upload_id = upload_file( + self._headers, self._url, file_path, self._verify + ) except requests.exceptions.HTTPError as http_err: return handle_http_error(self._logger, http_err) except requests.RequestException as e: @@ -451,6 +485,7 @@ def upload_and_create_from_openhtf_report( f"{self._url}/import", json=payload, headers=self._headers, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() @@ -478,6 +513,7 @@ def get_websocket_url(self) -> dict: response = requests.get( f"{self._url}/rooms", headers=self._headers, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) response.raise_for_status() diff --git a/tofupilot/openhtf/upload.py b/tofupilot/openhtf/upload.py index 42a3a93..56737b6 100644 --- a/tofupilot/openhtf/upload.py +++ b/tofupilot/openhtf/upload.py @@ -24,6 +24,14 @@ class upload: # pylint: disable=invalid-name This function behaves similarly to manually parsing the OpenHTF JSON test report and calling `TofuPilotClient().create_run()` with the parsed data. + Args: + api_key (Optional[str]): API key for authentication with TofuPilot's API. + allow_nan (Optional[bool]): Whether to allow NaN values in JSON serialization. + url (Optional[str]): Base URL for TofuPilot's API. + client (Optional[TofuPilotClient]): An existing TofuPilot client instance to use. + verify (Optional[str]): Path to a CA bundle file to verify TofuPilot's server certificate. + Useful for connecting to instances with custom/self-signed certificates. + ### Usage Example: ```python @@ -48,12 +56,14 @@ def __init__( allow_nan: Optional[bool] = False, url: Optional[str] = None, client: Optional[TofuPilotClient] = None, + verify: Optional[str] = None, ): self.allow_nan = allow_nan - self.client = client or TofuPilotClient(api_key=api_key, url=url) + self.client = client or TofuPilotClient(api_key=api_key, url=url, verify=verify) self._logger = self.client._logger self._url = self.client._url self._headers = self.client._headers + self._verify = verify self._max_attachments = self.client._max_attachments self._max_file_size = self.client._max_file_size @@ -133,6 +143,7 @@ def __call__(self, test_record: TestRecord): initialize_url, data=json.dumps(payload), headers=self._headers, + verify=self._verify, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -148,7 +159,13 @@ def __call__(self, test_record: TestRecord): timeout=SECONDS_BEFORE_TIMEOUT, ) - notify_server(self._headers, self._url, upload_id, run_id) + notify_server( + self._headers, + self._url, + upload_id, + run_id, + self._verify, + ) self._logger.success( "Attachment %s successfully uploaded and linked to run.", diff --git a/tofupilot/utils/files.py b/tofupilot/utils/files.py index e8d3cf1..5dd3616 100644 --- a/tofupilot/utils/files.py +++ b/tofupilot/utils/files.py @@ -46,8 +46,19 @@ def upload_file( headers: dict, url: str, file_path: str, + verify: Optional[str] = None, ) -> bool: - """Initializes an upload and stores file in it""" + """Initializes an upload and stores file in it + + Args: + headers (dict): Request headers including authorization + url (str): Base API URL + file_path (str): Path to the file to upload + verify (Optional[str]): Path to a CA bundle file to verify the server certificate + + Returns: + str: The ID of the created upload + """ # Upload initialization initialize_url = f"{url}/uploads/initialize" file_name = os.path.basename(file_path) @@ -58,6 +69,7 @@ def upload_file( data=json.dumps(payload), headers=headers, timeout=SECONDS_BEFORE_TIMEOUT, + verify=verify, ) response.raise_for_status() @@ -78,14 +90,32 @@ def upload_file( return upload_id -def notify_server(headers: dict, url: str, upload_id: str, run_id: str) -> bool: - """Tells TP server to sync upload with newly created run""" +def notify_server( + headers: dict, + url: str, + upload_id: str, + run_id: str, + verify: Optional[str] = None, +) -> bool: + """Tells TP server to sync upload with newly created run + + Args: + headers (dict): Request headers including authorization + url (str): Base API URL + upload_id (str): ID of the upload to link + run_id (str): ID of the run to link to + verify (Optional[str]): Path to a CA bundle file to verify the server certificate + + Returns: + bool: True if successful + """ sync_url = f"{url}/uploads/sync" sync_payload = {"upload_id": upload_id, "run_id": run_id} response = requests.post( sync_url, data=json.dumps(sync_payload), + verify=verify, headers=headers, timeout=SECONDS_BEFORE_TIMEOUT, ) @@ -99,13 +129,23 @@ def upload_attachments( url: str, paths: List[Dict[str, Optional[str]]], run_id: str, + verify: Optional[str] = None, ): - """Creates one upload per file and stores them into TofuPilot""" + """Creates one upload per file and stores them into TofuPilot + + Args: + logger (Logger): Logger instance + headers (dict): Request headers including authorization + url (str): Base API URL + paths (List[Dict[str, Optional[str]]]): List of file paths to upload + run_id (str): ID of the run to link files to + verify (Optional[str]): Path to a CA bundle file to verify the server certificate + """ for file_path in paths: logger.info("Uploading %s...", file_path) - upload_id = upload_file(headers, url, file_path) - notify_server(headers, url, upload_id, run_id) + upload_id = upload_file(headers, url, file_path, verify) + notify_server(headers, url, upload_id, run_id, verify) logger.success( f"Attachment {file_path} successfully uploaded and linked to run."