Skip to content

Commit

Permalink
feat: Add client for Artifacts API (#70)
Browse files Browse the repository at this point in the history
Signed-off-by: Vigneshwaran Selvaraj <vigneshwaran.selvaraj@ni.com>
  • Loading branch information
vignesh-ni authored Sep 9, 2024
1 parent b3fff6b commit 0ada3d7
Show file tree
Hide file tree
Showing 7 changed files with 172 additions and 0 deletions.
28 changes: 28 additions & 0 deletions examples/artifact/upload_and_download_artifacts.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
import io

from nisystemlink.clients.artifact import ArtifactClient
from nisystemlink.clients.core import HttpConfiguration


# Setup the server configuration to point to your instance of SystemLink Enterprise
server_configuration = HttpConfiguration(
server_uri="https://yourserver.yourcompany.com",
api_key="YourAPIKeyGeneratedFromSystemLink",
)
client = ArtifactClient(configuration=server_configuration)

# Define the workspace and artifact content
workspace = "your workspace ID"
artifact_stream = io.BytesIO(b"test content")

# Upload the artifact
upload_response = client.upload_artifact(workspace=workspace, artifact=artifact_stream)
if upload_response and upload_response.id:
print(f"Uploaded artifact ID: {upload_response.id}")

# Download the artifact using the ID from the upload response
artifact_id = upload_response.id
download_response = client.download_artifact(artifact_id)
if download_response:
downloaded_content = download_response.read()
print(f"Downloaded artifact content: {downloaded_content.decode('utf-8')}")
3 changes: 3 additions & 0 deletions nisystemlink/clients/artifact/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from ._artifact_client import ArtifactClient

# flake8: noqa
83 changes: 83 additions & 0 deletions nisystemlink/clients/artifact/_artifact_client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
"""Implementation of ArtifactClient"""

from typing import BinaryIO, Optional

from nisystemlink.clients import core
from nisystemlink.clients.core._uplink._base_client import BaseClient
from nisystemlink.clients.core._uplink._methods import get, post, response_handler
from nisystemlink.clients.core.helpers._iterator_file_like import IteratorFileLike
from requests.models import Response
from uplink import Part, Path

from . import models


class ArtifactClient(BaseClient):
def __init__(self, configuration: Optional[core.HttpConfiguration] = None):
"""Initialize an instance.
Args:
configuration: Defines the web server to connect to and information about
how to connect. If not provided, the
:class:`HttpConfigurationManager <nisystemlink.clients.core.HttpConfigurationManager>`
is used to obtain the configuration.
Raises:
ApiException: if unable to communicate with the Notebook execution Service.
"""
if configuration is None:
configuration = core.HttpConfigurationManager.get_configuration()

super().__init__(configuration, base_path="/ninbartifact/v1/")

@post("artifacts")
def __upload_artifact(
self, workspace: Part, artifact: Part
) -> models.UploadArtifactResponse:
"""Uploads an artifact using multipart/form-data headers to send the file payload in the HTTP body.
Args:
workspace: The workspace containing the artifact.
artifact: The artifact to upload.
Returns:
UploadArtifactResponse: The response containing the artifact ID.
"""

def upload_artifact(
self, workspace: str, artifact: BinaryIO
) -> models.UploadArtifactResponse:
"""Uploads an artifact.
Args:
workspace: The workspace containing the artifact.
artifact: The artifact to upload.
Returns:
UploadArtifactResponse: The response containing the artifact ID.
"""
response = self.__upload_artifact(
workspace=workspace,
artifact=artifact,
)

return response

def _iter_content_filelike_wrapper(response: Response) -> IteratorFileLike:
return IteratorFileLike(response.iter_content(chunk_size=4096))

@response_handler(_iter_content_filelike_wrapper)
@get("artifacts/{id}")
def download_artifact(self, id: Path) -> IteratorFileLike:
"""Downloads an artifact.
Args:
id: The ID of the artifact to download.
Returns:
A file-like object for reading the artifact content.
"""
...
3 changes: 3 additions & 0 deletions nisystemlink/clients/artifact/models/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from ._upload_artifact_response import UploadArtifactResponse

# flake8: noqa
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
from nisystemlink.clients.core._uplink._json_model import JsonModel


class UploadArtifactResponse(JsonModel):
"""Response for an artifact upload request."""

id: str
"""Information about the uploaded artifact."""
1 change: 1 addition & 0 deletions tests/integration/artifact/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
# flake8: noqa
46 changes: 46 additions & 0 deletions tests/integration/artifact/test_artifact.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
import io
import os

import pytest
from nisystemlink.clients.artifact import ArtifactClient
from nisystemlink.clients.core._http_configuration import HttpConfiguration


@pytest.fixture(scope="class")
def client(enterprise_config: HttpConfiguration) -> ArtifactClient:
"""Fixture to create an ArtifactClient instance."""
return ArtifactClient(enterprise_config)


@pytest.mark.integration
@pytest.mark.enterprise
class TestArtifact:

def test__upload_artifact__artifact_uploaded(self, client: ArtifactClient):
workspace = os.getenv("SYSTEMLINK_WORKSPACE_ID")

if workspace is not None:
artifact_stream = io.BytesIO(b"test content")

response = client.upload_artifact(
workspace=workspace, artifact=artifact_stream
)

assert response is not None
assert response.id is not None

def test__download_artifact__artifact_downloaded(self, client: ArtifactClient):
workspace = os.getenv("SYSTEMLINK_WORKSPACE_ID")

if workspace is not None:
artifact_content = b"test content"
artifact_stream = io.BytesIO(artifact_content)

upload_response = client.upload_artifact(
workspace=workspace, artifact=artifact_stream
)
artifact_id = upload_response.id
download_response = client.download_artifact(artifact_id)

assert download_response is not None
assert download_response.read() == artifact_content

0 comments on commit 0ada3d7

Please sign in to comment.