Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 41 additions & 5 deletions tofupilot/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)

Expand All @@ -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.",
Expand Down Expand Up @@ -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,
)
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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:
Expand All @@ -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()
Expand Down Expand Up @@ -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()
Expand Down
21 changes: 19 additions & 2 deletions tofupilot/openhtf/upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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,
)

Expand All @@ -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.",
Expand Down
52 changes: 46 additions & 6 deletions tofupilot/utils/files.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -58,6 +69,7 @@ def upload_file(
data=json.dumps(payload),
headers=headers,
timeout=SECONDS_BEFORE_TIMEOUT,
verify=verify,
)

response.raise_for_status()
Expand All @@ -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,
)
Expand All @@ -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."
Expand Down