Merge pull request #63 from tofupilot/alex/app-2850-enable-custom-certificate

feat: add possibility for custom certificate in requests

Author: upview

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,14 +174,20 @@ 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)
 
             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()

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.",

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."