def test_read_n(self): """Test that the stream only reads 'n' elements.""" self.mock_result.set_state(PipelineState.DONE) self.cache.write(list(range(5)), 'full', self.cache_key) self.cache.save_pcoder(None, 'full', self.cache_key) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=1, max_duration_secs=1) self.assertEqual(list(stream.read()), [0]) self.assertTrue(stream.is_done()) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=2, max_duration_secs=1) self.assertEqual(list(stream.read()), [0, 1]) self.assertTrue(stream.is_done()) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=5, max_duration_secs=1) self.assertEqual(list(stream.read()), list(range(5))) self.assertTrue(stream.is_done()) # Test that if the user asks for more than in the cache it still returns. stream = ElementStream( self.pcoll, '', self.cache_key, max_n=10, max_duration_secs=1) self.assertEqual(list(stream.read()), list(range(5))) self.assertTrue(stream.is_done())
def test_read(self): """Test reading and if a stream is done no more elements are returned.""" self.mock_result.set_state(PipelineState.DONE) self.cache.write(['expected'], 'full', self.cache_key) self.cache.save_pcoder(None, 'full', self.cache_key) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=1, max_duration_secs=1) self.assertFalse(stream.is_done()) self.assertEqual(list(stream.read())[0], 'expected') self.assertTrue(stream.is_done())
def test_done_if_terminated(self): """Test that terminating the job sets the stream as done.""" self.cache.write(['expected'], 'full', self.cache_key) self.cache.save_pcoder(None, 'full', self.cache_key) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=100, max_duration_secs=10) self.assertFalse(stream.is_done()) self.assertEqual(list(stream.read(tail=False))[0], 'expected') # The limiters were not reached, so the stream is not done yet. self.assertFalse(stream.is_done()) self.mock_result.set_state(PipelineState.DONE) self.assertEqual(list(stream.read(tail=False))[0], 'expected') # The underlying pipeline is terminated, so the stream won't yield new # elements. self.assertTrue(stream.is_done())