def test_get_generate_dataset_task_checks_for_existing_task(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) controller_id = bson.ObjectId() simulator_id = bson.ObjectId() simulator_config = { 'stereo_offset': 0.15, 'provide_rgb': True, 'provide_depth': True, 'provide_labels': False, 'provide_world_normals': False } repeat = 170 subject.get_generate_dataset_task(controller_id, simulator_id, simulator_config, repeat) self.assertTrue(mock_collection.find_one.called) query = mock_collection.find_one.call_args[0][0] self.assertIn('controller_id', query) self.assertEqual(controller_id, query['controller_id']) self.assertIn('simulator_id', query) self.assertEqual(simulator_id, query['simulator_id']) for key, value in simulator_config.items(): self.assertIn('simulator_config.{}'.format(key), query) self.assertEqual(value, query['simulator_config.{}'.format(key)]) self.assertIn('repeat', query) self.assertEqual(repeat, query['repeat'])
def test_get_import_dataset_task_returns_new_instance_if_no_existing(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_collection.find_one.return_value = None mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) module_name = 'test_module' path = '/tmp/dataset/thisisadataset' result = subject.get_import_dataset_task(module_name, path) self.assertIsInstance(result, import_dataset_task.ImportDatasetTask) self.assertIsNone(result.identifier)
def test_get_benchmark_returns_new_instance_if_no_existing(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_collection.find_one.return_value = None mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) trial_result_id = bson.ObjectId() benchmark_id = bson.ObjectId() result = subject.get_benchmark_task(trial_result_id, benchmark_id) self.assertIsInstance(result, benchmark_task.BenchmarkTrialTask) self.assertIsNone(result.identifier)
def test_get_run_system_task_returns_new_instance_if_no_existing(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_collection.find_one.return_value = None mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) system_id = bson.ObjectId() image_source_id = bson.ObjectId() result = subject.get_run_system_task(system_id, image_source_id) self.assertIsInstance(result, run_system_task.RunSystemTask) self.assertIsNone(result.identifier)
def test_get_generate_dataset_task_returns_new_instance_if_no_existing( self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_collection.find_one.return_value = None mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) result = subject.get_generate_dataset_task(bson.ObjectId(), bson.ObjectId(), {'provide_rgb': True}) self.assertIsInstance(result, generate_dataset_task.GenerateDatasetTask) self.assertIsNone(result.identifier)
def test_get_import_dataset_task_returns_deserialized_existing(self): s_task = {'_type': 'ImportDatasetTask', '_id': bson.ObjectId()} mock_entity = mock.MagicMock() mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_collection.find_one.return_value = s_task mock_db_client = mock.create_autospec(database.client.DatabaseClient) mock_db_client.deserialize_entity.return_value = mock_entity subject = manager.TaskManager(mock_collection, mock_db_client) result = subject.get_import_dataset_task('lol no', '/tmp/lolno') self.assertTrue(mock_db_client.deserialize_entity.called) self.assertEqual(s_task, mock_db_client.deserialize_entity.call_args[0][0]) self.assertEqual(mock_entity, result)
def test_get_benchmark_task_checks_for_existing_task(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) trial_result_id = bson.ObjectId() benchmark_id = bson.ObjectId() subject.get_benchmark_task(trial_result_id, benchmark_id) self.assertTrue(mock_collection.find_one.called) query = mock_collection.find_one.call_args[0][0] self.assertIn('trial_result_id', query) self.assertEqual(trial_result_id, query['trial_result_id']) self.assertIn('benchmark_id', query) self.assertEqual(benchmark_id, query['benchmark_id'])
def test_get_run_system_task_checks_for_existing_task(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) system_id = bson.ObjectId() image_source_id = bson.ObjectId() subject.get_run_system_task(system_id, image_source_id) self.assertTrue(mock_collection.find_one.called) query = mock_collection.find_one.call_args[0][0] self.assertIn('system_id', query) self.assertEqual(system_id, query['system_id']) self.assertIn('image_source_id', query) self.assertEqual(image_source_id, query['image_source_id'])
def test_do_task_checks_run_system_task_is_unique(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) system_id = bson.ObjectId() image_source_id = bson.ObjectId() task = run_system_task.RunSystemTask(system_id, image_source_id) subject.do_task(task) self.assertTrue(mock_collection.find.called) query = mock_collection.find.call_args[0][0] self.assertIn('system_id', query) self.assertEqual(system_id, query['system_id']) self.assertIn('image_source_id', query) self.assertEqual(image_source_id, query['image_source_id'])
def test_do_task_checks_train_system_task_is_unique(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) trainer_id = bson.ObjectId() trainee_id = bson.ObjectId() task = train_system_task.TrainSystemTask(trainer_id, trainee_id) subject.do_task(task) self.assertTrue(mock_collection.find.called) query = mock_collection.find.call_args[0][0] self.assertIn('trainer_id', query) self.assertEqual(trainer_id, query['trainer_id']) self.assertIn('trainee_id', query) self.assertEqual(trainee_id, query['trainee_id'])
def test_do_task_checks_import_benchmark_task_is_unique(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) module_name = 'test_module' path = '/tmp/dataset/thisisadataset' task = import_dataset_task.ImportDatasetTask(module_name, path) subject.do_task(task) self.assertTrue(mock_collection.find.called) query = mock_collection.find.call_args[0][0] self.assertIn('module_name', query) self.assertEqual(module_name, query['module_name']) self.assertIn('path', query) self.assertEqual(path, query['path'])
def test_get_import_dataset_task_checks_for_existing_task(self): mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) module_name = 'test_module' path = '/tmp/dataset/thisisadataset' additional_args = {'foo': 'bar'} subject.get_import_dataset_task(module_name, path, additional_args) self.assertTrue(mock_collection.find_one.called) query = mock_collection.find_one.call_args[0][0] self.assertIn('module_name', query) self.assertEqual(module_name, query['module_name']) self.assertIn('path', query) self.assertEqual(path, query['path']) self.assertIn('additional_args.foo', query) self.assertEqual('bar', query['additional_args.foo'])
def test_do_task_saves_new_task(self): # Mockthe method chain on the pymongo cursor mock_cursor = mock.MagicMock() mock_cursor.limit.return_value = mock_cursor mock_cursor.count.return_value = 0 mock_collection = mock.create_autospec(pymongo.collection.Collection) mock_collection.find.return_value = mock_cursor mock_db_client = mock.create_autospec(database.client.DatabaseClient) subject = manager.TaskManager(mock_collection, mock_db_client) system_id = bson.ObjectId() image_source_id = bson.ObjectId() task = run_system_task.RunSystemTask(system_id, image_source_id) subject.do_task(task) self.assertTrue(mock_collection.insert.called) s_task = task.serialize() del s_task[ '_id'] # This gets set after the insert call, clear it again self.assertEqual(s_task, mock_collection.insert.call_args[0][0])