def test_get_datasets_returns_only_datasets_with_requested_quality(
self, mock_load_sequences, mock_task_manager):
mock_task = mock.Mock()
mock_task.is_finished = False
mock_task_manager.get_import_dataset_task.return_value = mock_task
quality = ndds.QualityLevel.MIN_QUALITY
num_sequences = sum(1 for sequence_entry in self.sequence_entries
if sequence_entry.quality_level == quality)
mock_load_sequences.return_value = self.sequence_entries
manager = ndds.NDDSManager(self.root_path)
_, pending = manager.get_datasets(quality_level=quality)
self.assertEqual(num_sequences, pending)
self.assertEqual(num_sequences,
mock_task_manager.get_import_dataset_task.call_count)
for sequence_entry in self.sequence_entries:
if sequence_entry.quality_level == quality:
self.assertIn(
mock.call(module_name=ndds_loader.__name__,
path=str(sequence_entry.path),
additional_args={
'sequence_name': sequence_entry.sequence_name
},
num_cpus=mock.ANY,
num_gpus=mock.ANY,
memory_requirements=mock.ANY,
expected_duration=mock.ANY),
mock_task_manager.get_import_dataset_task.call_args_list)
def test_get_datasets_allows_multiple_conditions(self, mock_load_sequences,
mock_task_manager):
mock_task = mock.Mock()
mock_task.is_finished = False
mock_task_manager.get_import_dataset_task.return_value = mock_task
trajectory_id = self.sequence_entries[0].trajectory_id
time_of_day = TimeOfDay.DAY
num_sequences = sum(1 for sequence_entry in self.sequence_entries
if sequence_entry.time_of_day == time_of_day
and sequence_entry.trajectory_id == trajectory_id)
mock_load_sequences.return_value = self.sequence_entries
manager = ndds.NDDSManager(self.root_path)
_, pending = manager.get_datasets(time_of_day=time_of_day,
trajectory_id=trajectory_id)
self.assertEqual(num_sequences, pending)
self.assertEqual(num_sequences,
mock_task_manager.get_import_dataset_task.call_count)
for sequence_entry in self.sequence_entries:
if sequence_entry.time_of_day == time_of_day and sequence_entry.trajectory_id == trajectory_id:
self.assertIn(
mock.call(module_name=ndds_loader.__name__,
path=str(sequence_entry.path),
additional_args={
'sequence_name': sequence_entry.sequence_name
},
num_cpus=mock.ANY,
num_gpus=mock.ANY,
memory_requirements=mock.ANY,
expected_duration=mock.ANY),
mock_task_manager.get_import_dataset_task.call_args_list)
def test_get_datasets_returns_only_dataset_from_requested_environment(
self, mock_load_sequences, mock_task_manager):
mock_task = mock.Mock()
mock_task.is_finished = False
mock_task_manager.get_import_dataset_task.return_value = mock_task
environment = self.sequence_entries[0].environment
num_sequences = sum(1 for sequence_entry in self.sequence_entries
if sequence_entry.environment == environment)
mock_load_sequences.return_value = self.sequence_entries
manager = ndds.NDDSManager(self.root_path)
_, pending = manager.get_datasets(environment=environment)
self.assertEqual(num_sequences, pending)
self.assertEqual(num_sequences,
mock_task_manager.get_import_dataset_task.call_count)
for sequence_entry in self.sequence_entries:
if sequence_entry.environment == environment:
self.assertIn(
mock.call(module_name=ndds_loader.__name__,
path=str(sequence_entry.path),
additional_args={
'sequence_name': sequence_entry.sequence_name
},
num_cpus=mock.ANY,
num_gpus=mock.ANY,
memory_requirements=mock.ANY,
expected_duration=mock.ANY),
mock_task_manager.get_import_dataset_task.call_args_list)
def test_get_datasets_returns_result_ids_from_complete_tasks(
self, mock_load_sequences, mock_task_manager):
result_id = ObjectId()
finished_task = mock.Mock()
finished_task.is_finished = True
finished_task.get_result.return_value = result_id
unfinished_task = mock.Mock()
unfinished_task.is_finished = False
trajectory_id = self.sequence_entries[0].trajectory_id
num_sequences = sum(1 for sequence_entry in self.sequence_entries
if sequence_entry.trajectory_id == trajectory_id)
num_pending = sum(
1 for sequence_entry in self.sequence_entries
if sequence_entry.trajectory_id == trajectory_id
and not str(sequence_entry.path).endswith('-000.tar.gz'))
mock_load_sequences.return_value = self.sequence_entries
mock_task_manager.get_import_dataset_task.side_effect = \
lambda path, *args, **kwargs: finished_task if str(path).endswith('-000.tar.gz') else unfinished_task
manager = ndds.NDDSManager(self.root_path)
results, pending = manager.get_datasets(trajectory_id=trajectory_id)
self.assertEqual(num_sequences,
mock_task_manager.get_import_dataset_task.call_count)
self.assertEqual(num_pending, pending)
self.assertEqual(num_sequences - num_pending, len(results))
for result in results:
self.assertEqual(result_id, result)
for sequence_entry in self.sequence_entries:
if sequence_entry.trajectory_id == trajectory_id:
self.assertIn(
mock.call(module_name=ndds_loader.__name__,
path=str(sequence_entry.path),
additional_args={
'sequence_name': sequence_entry.sequence_name
},
num_cpus=mock.ANY,
num_gpus=mock.ANY,
memory_requirements=mock.ANY,
expected_duration=mock.ANY),
mock_task_manager.get_import_dataset_task.call_args_list)
def test_get_datasets_returns_all_datasets_by_default(
self, mock_load_sequences, mock_task_manager):
mock_task = mock.Mock()
mock_task.is_finished = False
mock_task_manager.get_import_dataset_task.return_value = mock_task
mock_load_sequences.return_value = self.sequence_entries
manager = ndds.NDDSManager(self.root_path)
_, pending = manager.get_datasets()
self.assertEqual(len(self.sequence_entries), pending)
self.assertEqual(len(self.sequence_entries),
mock_task_manager.get_import_dataset_task.call_count)
for sequence_entry in self.sequence_entries:
self.assertIn(
mock.call(module_name=ndds_loader.__name__,
path=str(sequence_entry.path),
additional_args={
'sequence_name': sequence_entry.sequence_name
},
num_cpus=mock.ANY,
num_gpus=mock.ANY,
memory_requirements=mock.ANY,
expected_duration=mock.ANY),
mock_task_manager.get_import_dataset_task.call_args_list)