def test__generate_segments(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": ""
},
],
}
list_segments = mocker.patch(
f"{dataset.__name__}.FusionDatasetClient._list_segments",
return_value=response_data,
)
segment_generator = ReturnGenerator(
self.fusion_dataset_client._generate_segments())
assert [segment.name for segment in segment_generator
] == [item["name"] for item in response_data["segments"]]
list_segments.assert_called_once_with(offset, limit)
assert segment_generator.value == response_data["totalCount"]
def test__generate_jobs(self, mocker, mock_list_jobs):
status, offset, limit = None, 0, 128
params = {
"jobType": "squashAndMerge",
"status": status,
"offset": offset,
"limit": limit,
}
open_api_do, response_data = mock_list_jobs(mocker)
job_generator = ReturnGenerator(
self.dataset_client.squash_and_merge._generate_jobs())
assert list(job_generator) == [
SquashAndMergeJob.from_response_body(
item,
dataset_id=self.dataset_client.squash_and_merge._dataset_id,
client=self.dataset_client.squash_and_merge._client,
job_updater=self.dataset_client.squash_and_merge._get_job,
draft_getter=self.dataset_client.squash_and_merge.
_draft_getter,
) for item in response_data["jobs"]
]
open_api_do.assert_called_once_with("GET",
"jobs",
self.dataset_client.dataset_id,
params=params)
assert job_generator.value == response_data["totalCount"]
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__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")