def test_with_seed_start(self, mock_ia_seed): augseq = mock.MagicMock() augseq.augment_batch_.return_value = "augmented_batch_" image = np.zeros((1, 1, 3), dtype=np.uint8) batch = UnnormalizedBatch(images=[image]) batch_idx = 1 seed_start = 10 multicore.Pool._WORKER_AUGSEQ = augseq multicore.Pool._WORKER_SEED_START = seed_start result = multicore._Pool_worker(batch_idx, batch) # expected seeds used seed = seed_start + batch_idx seed_global_expected = ( iarandom.SEED_MIN_VALUE + (seed - 10**9) % (iarandom.SEED_MAX_VALUE - iarandom.SEED_MIN_VALUE)) seed_local_expected = ( iarandom.SEED_MIN_VALUE + seed % (iarandom.SEED_MAX_VALUE - iarandom.SEED_MIN_VALUE)) assert result == "augmented_batch_" assert augseq.augment_batch_.call_count == 1 augseq.augment_batch_.assert_called_once_with(batch) mock_ia_seed.assert_called_once_with(seed_global_expected) augseq.seed_.assert_called_once_with(seed_local_expected)
def test_without_seed_start(self): augseq = mock.MagicMock() augseq.augment_batch_.return_value = "augmented_batch_" image = np.zeros((1, 1, 3), dtype=np.uint8) batch = UnnormalizedBatch(images=[image]) multicore.Pool._WORKER_AUGSEQ = augseq result = multicore._Pool_worker(1, batch) assert result == "augmented_batch_" assert augseq.augment_batch_.call_count == 1 augseq.augment_batch_.assert_called_once_with(batch)