Fix stopping Watch for logs and events with timeout (#330)

Author: tomplus

kubernetes_asyncio/watch/watch.py

@@ -174,16 +174,20 @@ async def next(self):
 
             line = line.decode('utf8')
 
+            # Special case for faster log streaming
+            if self.return_type == 'str':
+                if line == '':
+                    # end of log
+                    raise StopAsyncIteration
+                return line
+
             # Stop the iterator if K8s sends an empty response. This happens when
             # eg the supplied timeout has expired.
             if line == '':
                 if watch_forever:
                     self._reconnect()
                     continue
-
-            # Special case for faster log streaming
-            if self.return_type == 'str':
-                return line
+                raise StopAsyncIteration
 
             # retry 410 error only once
             try:

kubernetes_asyncio/watch/watch_test.py

@@ -74,7 +74,7 @@ async def test_watch_for_follow(self):
         fake_resp = AsyncMock()
         fake_resp.content.readline = AsyncMock()
         fake_resp.release = Mock()
-        side_effects = ['log_line_1', 'log_line_2']
+        side_effects = ['log_line_1', 'log_line_2', '']
         side_effects = [_.encode('utf8') for _ in side_effects]
         side_effects.extend([AssertionError('Should not have been called')])
         fake_resp.content.readline.side_effect = side_effects
@@ -84,16 +84,12 @@ async def test_watch_for_follow(self):
         fake_api.read_namespaced_pod_log.__doc__ = ':param follow:\n:type follow: bool\n:rtype: str'
 
         watch = kubernetes_asyncio.watch.Watch()
-        count = 1
+        logs = []
         async with watch:
             async for e in watch.stream(fake_api.read_namespaced_pod_log):
-                self.assertEqual("log_line_1", e)
-                # Stop the watch. This must not return the next event which would
-                # be an AssertionError exception.
-                count += 1
-                if count == len(side_effects) - 1:
-                    watch.stop()
+                logs.append(e)
 
+        self.assertListEqual(logs, ['log_line_1', 'log_line_2'])
         fake_api.read_namespaced_pod_log.assert_called_once_with(
             _preload_content=False, follow=True)
         fake_resp.release.assert_called_once_with()