diff --git a/watch/watch.py b/watch/watch.py index 7e7e2cb7ef..6f10e7b521 100644 --- a/watch/watch.py +++ b/watch/watch.py @@ -63,6 +63,7 @@ def __init__(self, return_type=None): self._raw_return_type = return_type self._stop = False self._api_client = client.ApiClient() + self.resource_version = 0 def stop(self): self._stop = True @@ -81,6 +82,8 @@ def unmarshal_event(self, data, return_type): if return_type: obj = SimpleNamespace(data=json.dumps(js['raw_object'])) js['object'] = self._api_client.deserialize(obj, return_type) + if hasattr(js['object'], 'metadata'): + self.resource_version = js['object'].metadata.resource_version return js def stream(self, func, *args, **kwargs): @@ -113,12 +116,19 @@ def stream(self, func, *args, **kwargs): return_type = self.get_return_type(func) kwargs['watch'] = True kwargs['_preload_content'] = False - resp = func(*args, **kwargs) - try: - for line in iter_resp_lines(resp): - yield self.unmarshal_event(line, return_type) - if self._stop: - break - finally: - resp.close() - resp.release_conn() + + timeouts = ('timeout_seconds' in kwargs) + while True: + resp = func(*args, **kwargs) + try: + for line in iter_resp_lines(resp): + yield self.unmarshal_event(line, return_type) + if self._stop: + break + finally: + kwargs['resource_version'] = self.resource_version + resp.close() + resp.release_conn() + + if timeouts or self._stop: + break diff --git a/watch/watch_test.py b/watch/watch_test.py index 64b5835fe4..73bcc9410f 100644 --- a/watch/watch_test.py +++ b/watch/watch_test.py @@ -85,6 +85,38 @@ def test_watch_stream_twice(self): fake_resp.close.assert_called_once() fake_resp.release_conn.assert_called_once() + def test_watch_stream_loop(self): + w = Watch(float) + + fake_resp = Mock() + fake_resp.close = Mock() + fake_resp.release_conn = Mock() + fake_resp.read_chunked = Mock( + return_value=['{"type": "ADDED", "object": 1}\n']) + + fake_api = Mock() + fake_api.get_namespaces = Mock(return_value=fake_resp) + fake_api.get_namespaces.__doc__ = ':return: V1NamespaceList' + + count = 0 + + # when timeout_seconds is set, auto-exist when timeout reaches + for e in w.stream(fake_api.get_namespaces, timeout_seconds=1): + count = count + 1 + self.assertEqual(count, 1) + + # when no timeout_seconds, only exist when w.stop() is called + for e in w.stream(fake_api.get_namespaces): + count = count + 1 + if count == 2: + w.stop() + + self.assertEqual(count, 2) + self.assertEqual(fake_api.get_namespaces.call_count, 2) + self.assertEqual(fake_resp.read_chunked.call_count, 2) + self.assertEqual(fake_resp.close.call_count, 2) + self.assertEqual(fake_resp.release_conn.call_count, 2) + def test_unmarshal_with_float_object(self): w = Watch() event = w.unmarshal_event('{"type": "ADDED", "object": 1}', 'float')