def test_state_remains_consistent(self): random = np.random.RandomState(144135) subject = task.Task(state=task.JobState.UNSTARTED) for idx in range(50): change = random.randint(0, 3 if idx > 30 else 2) if idx > 30 and change == 2: subject.mark_job_complete(bson.ObjectId()) elif change == 1: subject.mark_job_started('test', random.randint(0, 1000)) else: subject.mark_job_failed() # Make sure that the node id, job id, and result match the state if subject.is_unstarted: self.assertIsNone(subject.node_id) self.assertIsNone(subject.job_id) self.assertIsNone(subject.result) elif subject.is_finished: self.assertIsNone(subject.node_id) self.assertIsNone(subject.job_id) self.assertIsNotNone(subject.result) else: self.assertIsNotNone(subject.node_id) self.assertIsNotNone(subject.job_id) self.assertIsNone(subject.result) # Make sure we don't have an empty update, and we don't set and unset the same keys set_keys = set() unset_keys = set() if '$set' in subject._updates: self.assertNotEqual({}, subject._updates['$set']) set_keys = set(subject._updates['$set'].keys()) if '$unset' in subject._updates: self.assertNotEqual({}, subject._updates['$unset']) unset_keys = set(subject._updates['$unset'].keys()) self.assertEqual(set(), set_keys & unset_keys)
def test_mark_job_failed_doesnt_affect_unstarted_jobs(self): subject = task.Task(state=task.JobState.UNSTARTED) self.assertTrue(subject.is_unstarted) self.assertFalse(subject.is_finished) subject.mark_job_failed() self.assertTrue(subject.is_unstarted) self.assertFalse(subject.is_finished) self.assertEqual({}, subject._updates)
def test_mark_job_complete_doesnt_affect_unstarted_jobs(self): subject = task.Task(state=task.JobState.UNSTARTED) self.assertTrue(subject.is_unstarted) self.assertFalse(subject.is_finished) subject.mark_job_complete(bson.ObjectId()) self.assertTrue(subject.is_unstarted) self.assertFalse(subject.is_finished) self.assertEqual({}, subject._updates) self.assertIsNone(subject.result)
def test_mark_job_failed_changes_running_to_unstarted(self): subject = task.Task(state=task.JobState.RUNNING, node_id='test', job_id=5) self.assertFalse(subject.is_unstarted) subject.mark_job_failed() self.assertTrue(subject.is_unstarted) self.assertIsNone(subject.node_id) self.assertIsNone(subject.job_id)
def test_mark_job_complete_changes_running_to_finished(self): result_id = bson.ObjectId() subject = task.Task(state=task.JobState.RUNNING, node_id='test', job_id=5) self.assertFalse(subject.is_unstarted) subject.mark_job_complete(result_id) self.assertTrue(subject.is_finished) self.assertIsNone(subject.node_id) self.assertIsNone(subject.job_id) self.assertEqual(result_id, subject.result)
def test_mark_job_started_stores_updates(self): subject = task.Task(state=task.JobState.UNSTARTED, id_=bson.ObjectId()) mock_collection = mock.create_autospec(pymongo.collection.Collection) subject.mark_job_started('test', 12) subject.save_updates(mock_collection) self.assertTrue(mock_collection.update.called) self.assertEqual({'_id': subject.identifier}, mock_collection.update.call_args[0][0]) query = mock_collection.update.call_args[0][1] self.assertIn('$set', query) self.assertIn('node_id', query['$set']) self.assertIn('job_id', query['$set'])
def test_mark_job_started_doesnt_affect_finished_jobs(self): result_id = bson.ObjectId() subject = task.Task(state=task.JobState.DONE, result=result_id) self.assertFalse(subject.is_unstarted) self.assertTrue(subject.is_finished) subject.mark_job_started('test', 12) self.assertFalse(subject.is_unstarted) self.assertTrue(subject.is_finished) self.assertIsNone(subject.node_id) self.assertIsNone(subject.job_id) self.assertEqual(result_id, subject.result) self.assertEqual({}, subject._updates)
def test_mark_job_started_doesnt_affect_already_running_jobs(self): subject = task.Task(state=task.JobState.RUNNING, node_id='external', job_id=3) self.assertFalse(subject.is_unstarted) self.assertFalse(subject.is_finished) subject.mark_job_started('test', 12) self.assertFalse(subject.is_unstarted) self.assertFalse(subject.is_finished) self.assertEqual('external', subject.node_id) self.assertEqual(3, subject.job_id) self.assertEqual({}, subject._updates)
def test_mark_job_complete_stores_updates(self): result_id = bson.ObjectId() subject = task.Task(state=task.JobState.RUNNING, node_id='test', job_id=14, id_=bson.ObjectId()) mock_collection = mock.create_autospec(pymongo.collection.Collection) subject.mark_job_complete(result_id) subject.save_updates(mock_collection) self.assertTrue(mock_collection.update.called) self.assertEqual({'_id': subject.identifier}, mock_collection.update.call_args[0][0]) query = mock_collection.update.call_args[0][1] self.assertIn('$set', query) self.assertIn('result', query['$set']) self.assertIn('$unset', query) self.assertIn('node_id', query['$unset']) self.assertIn('job_id', query['$unset'])
def make_instance(self, *args, **kwargs): kwargs = du.defaults( kwargs, { 'state': task.JobState.RUNNING, 'num_cpus': np.random.randint(0, 1000), 'num_gpus': np.random.randint(0, 1000), 'memory_requirements': '{}MB'.format(np.random.randint(0, 50000)), 'expected_duration': '{0}:{1}:{2}'.format(np.random.randint(1000), np.random.randint(60), np.random.randint(60)), 'node_id': 'node-{}'.format(np.random.randint(10000)), 'job_id': np.random.randint(1000) }) return task.Task(*args, **kwargs)
def test_mark_job_started_changes_unstarted_to_running(self): subject = task.Task(state=task.JobState.UNSTARTED) self.assertTrue(subject.is_unstarted) subject.mark_job_started('test', 12) self.assertFalse(subject.is_unstarted)