Skip to content

Commit

Permalink
RSDK-4978 - return file id from upload (#443)
Browse files Browse the repository at this point in the history
  • Loading branch information
@cheukt
cheukt committed Sep 27, 2023
1 parent 5050a0e commit 1100894
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 17 deletions.
10 changes: 6 additions & 4 deletions src/viam/app/data_client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -438,7 +438,7 @@ async def binary_data_capture_upload(
method_parameters: Optional[Mapping[str, Any]] = None,
tags: Optional[List[str]] = None,
data_request_times: Optional[Tuple[datetime, datetime]] = None,
) -> None:
) -> str:
"""Upload binary sensor data.
Upload binary data collected on a robot through a specific component (e.g., a motor) along with the relevant metadata to
Expand Down Expand Up @@ -479,7 +479,8 @@ async def binary_data_capture_upload(
method_parameters=method_parameters,
tags=tags,
)
await self._data_capture_upload(metadata=metadata, sensor_contents=[sensor_contents])
response = await self._data_capture_upload(metadata=metadata, sensor_contents=[sensor_contents])
return response.file_id

async def tabular_data_capture_upload(
self,
Expand All @@ -491,7 +492,7 @@ async def tabular_data_capture_upload(
method_parameters: Optional[Mapping[str, Any]] = None,
tags: Optional[List[str]] = None,
data_request_times: Optional[List[Tuple[datetime, datetime]]] = None,
) -> None:
) -> str:
"""Upload tabular sensor data.
Upload tabular data collected on a robot through a specific component (e.g., a motor) along with the relevant metadata to
Expand Down Expand Up @@ -550,7 +551,8 @@ async def tabular_data_capture_upload(
method_parameters=method_parameters,
tags=tags,
)
await self._data_capture_upload(metadata=metadata, sensor_contents=sensor_contents)
response = await self._data_capture_upload(metadata=metadata, sensor_contents=sensor_contents)
return response.file_id

async def _data_capture_upload(self, metadata: UploadMetadata, sensor_contents: List[SensorData]) -> DataCaptureUploadResponse:
request = DataCaptureUploadRequest(metadata=metadata, sensor_contents=sensor_contents)
Expand Down
4 changes: 2 additions & 2 deletions tests/mocks/services.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@
CreateKeyRequest,
CreateKeyResponse,
GetRobotAPIKeysRequest,
GetRobotAPIKeysResponse
GetRobotAPIKeysResponse,
)
from viam.proto.app.data import (
AddBoundingBoxToImageByIDRequest,
Expand Down Expand Up @@ -711,7 +711,7 @@ async def DataCaptureUpload(self, stream: Stream[DataCaptureUploadRequest, DataC
assert request is not None
self.metadata = request.metadata
self.sensor_contents = request.sensor_contents
await stream.send_message(DataCaptureUploadResponse())
await stream.send_message(DataCaptureUploadResponse(file_id=self.file_upload_response))

async def FileUpload(self, stream: Stream[FileUploadRequest, FileUploadResponse]) -> None:
request_metadata = await stream.recv_message()
Expand Down
21 changes: 10 additions & 11 deletions tests/test_data_sync_client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,7 @@
from grpclib.testing import ChannelFor

from viam.app.data_client import DataClient
from viam.proto.app.datasync import (
UploadMetadata,
SensorData
)
from viam.proto.app.datasync import UploadMetadata, SensorData
from viam.utils import struct_to_dict, datetime_to_timestamp

from .mocks.services import MockDataSync
Expand All @@ -21,7 +18,7 @@
SECONDS_END = 1689256810
NANOS_END = 10
TAGS = ["tag"]
BINARY_DATA = b'binary_data'
BINARY_DATA = b"binary_data"
METHOD_NAME = "method_name"
DATETIMES = (datetime.now(), datetime.now())
TIMESTAMPS = (datetime_to_timestamp(DATETIMES[0]), datetime_to_timestamp(DATETIMES[1]))
Expand All @@ -45,35 +42,37 @@ class TestClient:
async def test_binary_data_capture_upload(self, service: MockDataSync):
async with ChannelFor([service]) as channel:
client = DataClient(channel, DATA_SERVICE_METADATA)
await client.binary_data_capture_upload(
file_id = await client.binary_data_capture_upload(
part_id=PART_ID,
component_type=COMPONENT_TYPE,
component_name=COMPONENT_NAME,
method_name=METHOD_NAME,
method_parameters=METHOD_PARAMETERS,
tags=TAGS,
data_request_times=DATETIMES,
binary_data=BINARY_DATA
binary_data=BINARY_DATA,
)
self.assert_sensor_contents(sensor_contents=list(service.sensor_contents), is_binary=True)
self.assert_metadata(metadata=service.metadata)
assert file_id == FILE_UPLOAD_RESPONSE

@pytest.mark.asyncio
async def test_tabular_data_capture_upload(self, service: MockDataSync):
async with ChannelFor([service]) as channel:
client = DataClient(channel, DATA_SERVICE_METADATA)
await client.tabular_data_capture_upload(
file_id = await client.tabular_data_capture_upload(
part_id=PART_ID,
component_type=COMPONENT_TYPE,
component_name=COMPONENT_NAME,
method_name=METHOD_NAME,
method_parameters=METHOD_PARAMETERS,
tags=TAGS,
data_request_times=[DATETIMES],
tabular_data=cast(List[Mapping[str, Any]], TABULAR_DATA)
tabular_data=cast(List[Mapping[str, Any]], TABULAR_DATA),
)
self.assert_sensor_contents(sensor_contents=list(service.sensor_contents), is_binary=False)
self.assert_metadata(metadata=service.metadata)
assert file_id == FILE_UPLOAD_RESPONSE

@pytest.mark.asyncio
async def test_file_upload(self, service: MockDataSync):
Expand All @@ -88,7 +87,7 @@ async def test_file_upload(self, service: MockDataSync):
method_parameters=METHOD_PARAMETERS,
file_extension=FILE_EXT,
tags=TAGS,
data=BINARY_DATA
data=BINARY_DATA,
)
assert file_id == FILE_UPLOAD_RESPONSE
self.assert_metadata(service.metadata)
Expand All @@ -109,7 +108,7 @@ async def test_file_upload_from_path(self, service: MockDataSync, tmp_path):
method_name=METHOD_NAME,
method_parameters=METHOD_PARAMETERS,
tags=TAGS,
filepath=path.resolve()
filepath=path.resolve(),
)
assert file_id == FILE_UPLOAD_RESPONSE
self.assert_metadata(service.metadata)
Expand Down

0 comments on commit 1100894

Please sign in to comment.