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)
