def test_get_notes(self, mocker): params = self.dataset_client._status.get_status_info() response_data = { "isContinuous": True, "binPointCloudFields": ["X", "Y", "Z", "Intensity", "Ring"], } open_api_do = mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data), ) assert self.dataset_client.get_notes() == Notes.loads(response_data) open_api_do.assert_called_once_with( "GET", "notes", self.dataset_client._dataset_id, params=params )
def test__list_segment(self, mocker): params = self.dataset_client._status.get_status_info() offset, limit = 0, 128 params["offset"] = offset params["limit"] = limit open_api_do = mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(), ) self.dataset_client._list_segments() open_api_do.assert_called_once_with("GET", "segments", self.dataset_client._dataset_id, params=params)
def test_get_auth_storage_config(self, mocker): with pytest.raises(TypeError): self.gas_client.get_auth_storage_config("") mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data={ "configs": [], "totalCount": 0 }), ) with pytest.raises(ResourceNotExistError): self.gas_client.get_auth_storage_config("cloud_config") response_data = { "configs": [ { "name": "cloud_config", "type": "azure", "filePath": "tensorbay", "isGravitiStorage": True, }, ], "offset": 0, "recordSize": 1, "totalCount": 1, } mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data), ) assert self.gas_client.get_auth_storage_config( response_data["configs"][0]["name"]) == StorageConfig.loads( response_data["configs"][0])
def mock_create_commit(mocker): """Mock the createCommit OpenAPI. Arguments: mocker: The mocker fixture. Returns: The patched mocker and response data. """ response_data = {"commitId": 1} return ( mocker.patch(f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data)), response_data, )
def mock_close_draft(mocker): """Mock the closeDraft OpenAPI. Arguments: mocker: The mocker fixture. Returns: The patched mocker and response data. """ response_data = {"status": "CLOSED"} return ( mocker.patch(f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data)), response_data, )
def mock_update_draft(mocker): """Mock the updateDraft OpenAPI. Arguments: mocker: The mocker fixture. Returns: The patched mocker and response data. """ response_data = {"draftNumber": 1} return ( mocker.patch(f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data)), response_data, )
def test_list_segment(self, mocker): response_data = { "offset": 0, "recordSize": 2, "totalCount": 2, "segments": [ {"name": "test", "description": ""}, {"name": "train", "description": ""}, ], } mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data), ) assert list(self.dataset_client.list_segment_names()) == [ segment["name"] for segment in response_data["segments"] ]
def mock_create_job(mocker): """Mock the createJob OpenAPI. Arguments: mocker: The mocker fixture. Returns: The patched mocker and response data. """ response_data = {"jobId": "123", "createdAt": 1} return ( mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data), ), response_data, )
def mock_get_total_size(mocker, large: bool = False): """Mock the getTotalSize openAPI. Arguments: mocker: The mocker fixture. large: Whether the dataset size is large than free storage. Returns: The patched mocker and response data. """ response_data = {"totalSize": 7 if not large else sys.maxsize} return ( mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data), ), response_data, )
def test_upload_catalog(self, mocker): self.dataset_client._status.checkout(draft_number=1) open_api_do = mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(), ) catalog = Catalog.loads( {"CLASSIFICATION": { "categories": [{ "name": "cat" }] }}) self.dataset_client.upload_catalog(catalog) put_data = {"catalog": catalog.dumps()} put_data.update(self.dataset_client._status.get_status_info()) open_api_do.assert_called_once_with("PUT", "labels/catalogs", self.dataset_client._dataset_id, json=put_data)
def test__move_segment(self, mocker): self.dataset_client._status.checkout(draft_number=1) source_name, target_name = "default", "train" with pytest.raises(InvalidParamsError): self.dataset_client._move_segment(source_name, target_name, strategy="move") open_api_do = mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(), ) post_data = { "strategy": "abort", "source": {"segmentName": source_name}, "segmentName": target_name, } post_data.update(self.dataset_client._status.get_status_info()) self.dataset_client._move_segment(source_name, target_name) open_api_do.assert_called_once_with( "POST", "segments?move", self.dataset_client._dataset_id, json=post_data )
def test_get_catelog(self, mocker): params = self.dataset_client._status.get_status_info() response_data = { "catalog": { "CLASSIFICATION": { "categories": [{ "name": "cat" }] } } } open_api_do = mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data), ) assert self.dataset_client.get_catalog() == Catalog.loads( response_data["catalog"]) open_api_do.assert_called_once_with("GET", "labels/catalogs", self.dataset_client._dataset_id, params=params)
def test__generate_data_diffs(self, mocker): params = {"offset": 0, "limit": 128} basehead = "commit-fc2cf9f910b7446aaccc7b833fe4e178" segment_name = "test_segment" data_diffs = { "data": [ { "remotePath": "1629800110640323954/3.png", "action": "add", "file": {"action": "add"}, "label": {"action": ""}, }, { "remotePath": "1629800110640323954/4.png", "action": "add", "file": {"action": "add"}, "label": {"action": ""}, }, ], "totalCount": 2, } open_api_do = mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=data_diffs), ) data_diff_generator = ReturnGenerator( self.dataset_client._generate_data_diffs(basehead, segment_name) ) for data_diff_instance, data_diff in zip(data_diff_generator, data_diffs["data"]): assert data_diff_instance.action == data_diff["action"] assert data_diff_instance.label.action == data_diff["label"]["action"] assert data_diff_instance.file.action == data_diff["file"]["action"] assert data_diff_instance.remote_path == data_diff["remotePath"] assert data_diff_generator.value == data_diffs["totalCount"] open_api_do.assert_called_once_with( "GET", f"diffs/{basehead}/segments/{segment_name}/data", self.dataset_client._dataset_id, params=params, )
def test_squash_and_merge(self, mocker): post_data = { "title": "squash_merge-1", "sourceBranchName": "branch-1", "targetBranchName": "branch-2", "strategy": "abort", } response_data = {"draftNumber": 2} open_api_do = mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data), ) draft_number = self.dataset_client.squash_and_merge( "squash_merge-1", source_branch_name="branch-1", target_branch_name="branch-2", strategy="abort", ) open_api_do.assert_called_once_with("POST", "squashAndMerge", self.dataset_client.dataset_id, json=post_data) assert draft_number == 2
def test_list_auth_data(self, mocker): params = { "prefix": "cloud_path", "limit": 128, } response_data = { "cloudFiles": [ "cloud_path/file1.jpg", "cloud_path/file2.jpg", ], "truncated": 0, "nextMarker": "", } open_api_do = mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data), ) for auth_data, cloud_path in zip_longest( self.cloud_client.list_auth_data(params["prefix"]), response_data["cloudFiles"]): assert isinstance(auth_data, AuthData) assert auth_data.path == cloud_path open_api_do.assert_called_once_with( "GET", f"cloud/{self.cloud_client._name}/files", params=params)
def test_list_auth_storage_configs(self, mocker): response_data = { "configs": [ { "name": "cloud_config", "type": "azure", "accountName": "tensorbay", "containerName": "tensorbay", }, ], "offset": 0, "recordSize": 1, "totalCount": 1, } mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data), ) configs = self.gas_client.list_auth_storage_configs() assert list(configs) == response_data["configs"]
def test__generate_segment_names(self, mocker): params = self.dataset_client._status.get_status_info() offset, limit = 0, 128 params["offset"] = offset params["limit"] = limit response_data = { "offset": 0, "recordSize": 2, "totalCount": 2, "segments": [ {"name": "test", "description": ""}, {"name": "train", "description": ""}, ], } open_api_do = mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=response_data), ) assert list(self.dataset_client._generate_segment_names()) == [ segment["name"] for segment in response_data["segments"] ] open_api_do.assert_called_once_with( "GET", "segments", self.dataset_client._dataset_id, params=params )
def test__generate_segment_diffs(self, mocker): params = {"offset": 0, "limit": 128} basehead = "commit-fc2cf9f910b7446aaccc7b833fe4e178" segment_diffs = { "segments": [{ "name": "test_segment", "action": "modify", "data": { "stats": { "total": 2, "additions": 2, "deletions": 0, "modifications": 0 } }, "sensor": { "action": "" }, }], "totalCount": 1, } open_api_do = mocker.patch( f"{gas.__name__}.Client.open_api_do", return_value=mock_response(data=segment_diffs), ) data_diffs = [ DataDiff.loads({ "remotePath": "1629800110640323954/3.png", "action": "add", "file": { "action": "add" }, "label": { "action": "" }, }), DataDiff.loads({ "remotePath": "1629800110640323954/4.png", "action": "add", "file": { "action": "add" }, "label": { "action": "" }, }), ] list_data_diffs = mocker.patch( f"{dataset.__name__}.DatasetClient._list_data_diffs", return_value=data_diffs, ) segment_diff_generator = ReturnGenerator( self.dataset_client._generate_segment_diffs(basehead)) for segment_diff_instance, segment_diff in zip( segment_diff_generator, segment_diffs["segments"]): assert len(segment_diff_instance ) == segment_diff["data"]["stats"]["total"] assert segment_diff_instance.action == segment_diff["action"] assert segment_diff_instance.name == "test_segment" for data_diff_instance, data_diff in zip(segment_diff_instance, data_diffs): assert data_diff_instance.action == data_diff.action assert data_diff_instance.label.action == data_diff.label.action assert data_diff_instance.file.action == data_diff.file.action assert data_diff_instance.remote_path == data_diff.remote_path assert segment_diff_generator.value == segment_diffs["totalCount"] open_api_do.assert_called_once_with("GET", f"diffs/{basehead}/segments", self.dataset_client._dataset_id, params=params) list_data_diffs.assert_called_once_with(basehead, "test_segment")
def test_enable_cache( self, mocker, mock_list_segments, mock_list_data_details, mock_get_total_size, tmp_path, caplog, ): self.dataset_client._status.checkout(draft_number=1) with pytest.raises(StatusError): self.dataset_client.enable_cache() self.dataset_client._status.checkout(commit_id="commit-1") _, response_data = mock_get_total_size(mocker, True) _, _, free_storage = shutil.disk_usage(tempfile.gettempdir()) self.dataset_client.enable_cache() log_message = ( "bytes left on device, " f'less than the dataset size {response_data["totalSize"]} bytes.\n ' "Please be aware that there is not enough space to cache the entire dataset." ) assert log_message in caplog.text mock_get_total_size(mocker) cache_path = tmp_path / "cache_data" dataset_cache_path = cache_path / self.dataset_client.dataset_id self.dataset_client.enable_cache(cache_path) assert self.dataset_client.cache_enabled == True assert self.dataset_client._cache_path == str(dataset_cache_path) _, list_segments = mock_list_segments(mocker) segment_name = list_segments["segments"][0]["name"] segment = self.dataset_client.get_segment(segment_name) mock_list_data_details(mocker) segment_data = segment.list_data() urlopen = mocker.patch( f"{file.__name__}.RemoteFileMixin._urlopen", return_value=mock_response(read=lambda *args: bytes(1)), ) segment_cache_path = (dataset_cache_path / self.dataset_client.status.commit_id / segment_name) # Traverse the segment for the first time. for data in segment_data: assert data.cache_path == str(segment_cache_path / data.path) data.open() local_open = mocker.patch( f"builtins.open", return_value=bytes(1), ) segment_length = len(segment_data) # Traverse the segment for 3 times using cached data. epoch = 3 for i in range(epoch): for data in segment_data: assert data.cache_path == str(segment_cache_path / data.path) data.open() # local_open.assert_called_once() assert local_open.call_count == segment_length * epoch assert urlopen.call_count == segment_length self.dataset_client._status.checkout(draft_number=1) assert self.dataset_client.cache_enabled == False assert self.dataset_client._cache_path == str(dataset_cache_path)