Merge pull request #93 from mitar/watch_and_logs

Making watch work with read_namespaced_pod_log

Author: k8s-ci-robot

watch/watch.py

@@ -20,6 +20,7 @@
 from kubernetes import client
 
 PYDOC_RETURN_LABEL = ":return:"
+PYDOC_FOLLOW_PARAM = ":param bool follow:"
 
 # Removing this suffix from return type name should give us event's object
 # type. e.g., if list_namespaces() returns "NamespaceList" type,
@@ -65,7 +66,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
+        self.resource_version = None
 
     def stop(self):
         self._stop = True
@@ -78,8 +79,17 @@ def get_return_type(self, func):
             return return_type[:-len(TYPE_LIST_SUFFIX)]
         return return_type
 
+    def get_watch_argument_name(self, func):
+        if PYDOC_FOLLOW_PARAM in pydoc.getdoc(func):
+            return 'follow'
+        else:
+            return 'watch'
+
     def unmarshal_event(self, data, return_type):
-        js = json.loads(data)
+        try:
+            js = json.loads(data)
+        except ValueError:
+            return data
         js['raw_object'] = js['object']
         if return_type:
             obj = SimpleNamespace(data=json.dumps(js['raw_object']))
@@ -122,7 +132,7 @@ def stream(self, func, *args, **kwargs):
 
         self._stop = False
         return_type = self.get_return_type(func)
-        kwargs['watch'] = True
+        kwargs[self.get_watch_argument_name(func)] = True
         kwargs['_preload_content'] = False
         if 'resource_version' in kwargs:
             self.resource_version = kwargs['resource_version']
@@ -136,9 +146,12 @@ def stream(self, func, *args, **kwargs):
                     if self._stop:
                         break
             finally:
-                kwargs['resource_version'] = self.resource_version
                 resp.close()
                 resp.release_conn()
+                if self.resource_version is not None:
+                    kwargs['resource_version'] = self.resource_version
+                else:
+                    self._stop = True
 
             if timeouts or self._stop:
                 break

watch/watch_test.py

@@ -67,6 +67,35 @@ def test_watch_with_decode(self):
         fake_resp.close.assert_called_once()
         fake_resp.release_conn.assert_called_once()
 
+    def test_watch_for_follow(self):
+        fake_resp = Mock()
+        fake_resp.close = Mock()
+        fake_resp.release_conn = Mock()
+        fake_resp.read_chunked = Mock(
+            return_value=[
+                'log_line_1\n',
+                'log_line_2\n'])
+
+        fake_api = Mock()
+        fake_api.read_namespaced_pod_log = Mock(return_value=fake_resp)
+        fake_api.read_namespaced_pod_log.__doc__ = ':param bool follow:\n:return: str'
+
+        w = Watch()
+        count = 1
+        for e in w.stream(fake_api.read_namespaced_pod_log):
+            self.assertEqual("log_line_1", e)
+            count += 1
+            # make sure we can stop the watch and the last event with won't be
+            # returned
+            if count == 2:
+                w.stop()
+
+        fake_api.read_namespaced_pod_log.assert_called_once_with(
+            _preload_content=False, follow=True)
+        fake_resp.read_chunked.assert_called_once_with(decode_content=False)
+        fake_resp.close.assert_called_once()
+        fake_resp.release_conn.assert_called_once()
+
     def test_watch_resource_version_set(self):
         # https://github.com/kubernetes-client/python/issues/700
         # ensure watching from a resource version does reset to resource