From 1100894dd30734936dc713d568b9e2057b24e18f Mon Sep 17 00:00:00 2001 From: Cheuk <90270663+cheukt@users.noreply.github.com> Date: Wed, 27 Sep 2023 13:49:49 -0400 Subject: [PATCH] RSDK-4978 - return file id from upload (#443) --- src/viam/app/data_client.py | 10 ++++++---- tests/mocks/services.py | 4 ++-- tests/test_data_sync_client.py | 21 ++++++++++----------- 3 files changed, 18 insertions(+), 17 deletions(-) diff --git a/src/viam/app/data_client.py b/src/viam/app/data_client.py index 703a88881..69dc70a46 100644 --- a/src/viam/app/data_client.py +++ b/src/viam/app/data_client.py @@ -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 @@ -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, @@ -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 @@ -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) diff --git a/tests/mocks/services.py b/tests/mocks/services.py index 0b57bee36..7ea6a9eb9 100644 --- a/tests/mocks/services.py +++ b/tests/mocks/services.py @@ -138,7 +138,7 @@ CreateKeyRequest, CreateKeyResponse, GetRobotAPIKeysRequest, - GetRobotAPIKeysResponse + GetRobotAPIKeysResponse, ) from viam.proto.app.data import ( AddBoundingBoxToImageByIDRequest, @@ -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() diff --git a/tests/test_data_sync_client.py b/tests/test_data_sync_client.py index 3b4e4814d..1a6e8ed3d 100644 --- a/tests/test_data_sync_client.py +++ b/tests/test_data_sync_client.py @@ -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 @@ -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])) @@ -45,7 +42,7 @@ 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, @@ -53,16 +50,17 @@ async def test_binary_data_capture_upload(self, service: MockDataSync): 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, @@ -70,10 +68,11 @@ async def test_tabular_data_capture_upload(self, service: MockDataSync): 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): @@ -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) @@ -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)