def test_list_tensors_with_keyerror(self):
"""Test list_tensors method with parameter tag raises keyerror."""
summary_base_dir = tempfile.mkdtemp()
train_job_01 = 'train_01'
name_01 = 'train_job_01'
log_path_01 = os.path.join(summary_base_dir, 'dir1')
self._make_path_and_file_list(log_path_01)
modify_time_01 = 1575460551.9777446
ms_loader = MSDataLoader(log_path_01)
loader_01 = DataLoader(log_path_01)
loader_01._loader = ms_loader
loader = LoaderStruct(loader_id=train_job_01,
name=name_01,
path=log_path_01,
latest_update_time=modify_time_01,
data_loader=loader_01)
loader_pool = {train_job_01: loader}
d_manager = DataManager([DataLoaderGenerator(summary_base_dir)])
d_manager._status = DataManagerStatus.LOADING.value
d_manager._loader_pool = loader_pool
tag = 'image'
with pytest.raises(ParamValueError):
d_manager.list_tensors(train_job_01, tag)
shutil.rmtree(summary_base_dir)
def test_list_tensors_success(self):
"""Test list_tensors method success."""
summary_base_dir = tempfile.mkdtemp()
train_job_01 = 'train_01'
name_01 = 'train_job_01'
log_path_01 = os.path.join(summary_base_dir, 'dir1')
self._make_path_and_file_list(log_path_01)
modify_time_01 = 1575460551.9777446
loader_01 = DataLoader(log_path_01)
ms_loader = MSDataLoader(log_path_01)
event_data = EventsData()
mock_obj = mock.MagicMock()
mock_obj.samples.return_value = {'test result'}
tag = 'image'
event_data._reservoir_by_tag = {tag: mock_obj}
ms_loader._events_data = event_data
loader_01._loader = ms_loader
loader = LoaderStruct(loader_id=train_job_01,
name=name_01,
path=log_path_01,
latest_update_time=modify_time_01,
data_loader=loader_01)
loader_pool = {train_job_01: loader}
d_manager = DataManager([DataLoaderGenerator(summary_base_dir)])
d_manager._status = DataManagerStatus.LOADING.value
d_manager._loader_pool = loader_pool
res = d_manager.list_tensors(train_job_01, tag)
assert res == {'test result'}
shutil.rmtree(summary_base_dir)