def test_clear(self):
cache = StreamingCache(cache_dir=None)
self.assertFalse(cache.exists('my_label'))
cache.write([TestStreamFileRecord()], 'my_label')
self.assertTrue(cache.exists('my_label'))
self.assertTrue(cache.clear('my_label'))
self.assertFalse(cache.exists('my_label'))
def test_exists(self):
cache = StreamingCache(cache_dir=None)
self.assertFalse(cache.exists('my_label'))
cache.write([TestStreamFileRecord()], 'my_label')
self.assertTrue(cache.exists('my_label'))
# '' shouldn't be treated as a wildcard to match everything.
self.assertFalse(cache.exists(''))
def test_clear(self):
cache = StreamingCache(cache_dir=None)
self.assertFalse(cache.exists('my_label'))
cache.sink(['my_label'], is_capture=True)
cache.write([TestStreamFileRecord()], 'my_label')
self.assertTrue(cache.exists('my_label'))
self.assertEqual(cache.capture_keys, set(['my_label']))
self.assertTrue(cache.clear('my_label'))
self.assertFalse(cache.exists('my_label'))
self.assertFalse(cache.capture_keys)
def test_capture_size_limit_not_reached_when_no_file(self):
cache = StreamingCache(cache_dir=None)
self.assertFalse(cache.exists('my_label'))
ie.current_env().set_cache_manager(cache, 'dummy pipeline')
limiter = capture_limiters.SizeLimiter(1)
self.assertFalse(limiter.is_triggered())
def test_empty(self):
CACHED_PCOLLECTION_KEY = repr(CacheKey('arbitrary_key', '', '', ''))
cache = StreamingCache(cache_dir=None)
self.assertFalse(cache.exists(CACHED_PCOLLECTION_KEY))
cache.write([], CACHED_PCOLLECTION_KEY)
reader, _ = cache.read(CACHED_PCOLLECTION_KEY)
# Assert that an empty reader returns an empty list.
self.assertFalse([e for e in reader])
def test_capture_size_limit_not_reached_when_file_size_under_limit(self):
ib.options.capture_size_limit = 100
cache = StreamingCache(cache_dir=None)
# Build a sink object to track the label as a capture in the test.
cache.sink(['my_label'], is_capture=True)
cache.write([TestStreamFileRecord()], 'my_label')
self.assertTrue(cache.exists('my_label'))
ie.current_env().set_cache_manager(cache)
self.assertFalse(ie.current_env().options.capture_control.
is_capture_size_limit_reached())
def test_capture_size_limit_not_reached_when_file_size_under_limit(self):
ib.options.capture_size_limit = 100
cache = StreamingCache(cache_dir=None)
# Build a sink object to track the label as a capture in the test.
cache.sink(['my_label'], is_capture=True)
cache.write([beam_interactive_api_pb2.TestStreamFileRecord()], 'my_label')
self.assertTrue(cache.exists('my_label'))
ie.current_env().set_cache_manager(cache, 'dummy pipeline')
limiter = capture_limiters.SizeLimiter(ib.options.capture_size_limit)
self.assertFalse(limiter.is_triggered())
def test_capture_size_limit_reached_when_file_size_above_limit(self):
ib.options.capture_size_limit = 1
cache = StreamingCache(cache_dir=None)
cache.sink(['my_label'], is_capture=True)
cache.write([
TestStreamFileRecord(recorded_event=TestStreamPayload.Event(
element_event=TestStreamPayload.Event.AddElements(elements=[
TestStreamPayload.TimestampedElement(
encoded_element=coders.FastPrimitivesCoder().encode(
'a'),
timestamp=0)
])))
], 'my_label')
self.assertTrue(cache.exists('my_label'))
ie.current_env().set_cache_manager(cache)
self.assertTrue(ie.current_env().options.capture_control.
is_capture_size_limit_reached())
def test_capture_size_limit_reached_when_file_size_above_limit(self):
ib.options.capture_size_limit = 1
cache = StreamingCache(cache_dir=None)
cache.sink(['my_label'], is_capture=True)
cache.write([
TestStreamFileRecord(recorded_event=TestStreamPayload.Event(
element_event=TestStreamPayload.Event.AddElements(elements=[
TestStreamPayload.TimestampedElement(
encoded_element=coders.FastPrimitivesCoder().encode(
'a'),
timestamp=0)
])))
], 'my_label')
self.assertTrue(cache.exists('my_label'))
p = _build_an_empty_streaming_pipeline()
ie.current_env().set_cache_manager(cache, p)
limiter = capture_limiters.SizeLimiter(1)
self.assertTrue(limiter.is_triggered())
def test_capture_size_limit_not_reached_when_no_file(self):
cache = StreamingCache(cache_dir=None)
self.assertFalse(cache.exists('my_label'))
ie.current_env().set_cache_manager(cache)
self.assertFalse(ie.current_env().options.capture_control.
is_capture_size_limit_reached())
