Repository: tomplus/kubernetes_asyncio (python)

Issue:
Calling watch.stop() does not close the response correctly
https://github.com/tomplus/kubernetes_asyncio/blob/d57e9e9be11f6789e1ce8d5b161acb64d29acf35/kubernetes_asyncio/watch/watch.py#L123-L124

It just raises an exception to exit.

I think we should add async context manager support for `Watch` objects, so it closes the response in `__aexit__`.

Pull request description:
fix #60
This PR adds async context manager support on `Watch` object, so the response is correctly closed when watch is stopped with `stop()` or by some exceptions.

fix #60 

Patch:
diff --git a/examples/example2.py b/examples/example2.py
index fc6b0742..ddbe25da 100644
--- a/examples/example2.py
+++ b/examples/example2.py
@@ -20,6 +20,9 @@ async def main():
             w.stop()
 
     print("Ended.")
+    # An explicit close is necessary to stop the stream
+    # or use async context manager like in example4.py
+    w.close()
 
 
 if __name__ == '__main__':
diff --git a/examples/example4.py b/examples/example4.py
index 6c63c20e..95de23bb 100644
--- a/examples/example4.py
+++ b/examples/example4.py
@@ -6,16 +6,18 @@ from kubernetes_asyncio import client, config, watch
 
 async def watch_namespaces():
     v1 = client.CoreV1Api()
-    async for event in watch.Watch().stream(v1.list_namespace):
-        etype, obj = event['type'], event['object']
-        print("{} namespace {}".format(etype, obj.metadata.name))
+    async with watch.Watch().stream(v1.list_namespace) as stream:
+        async for event in stream:
+            etype, obj = event['type'], event['object']
+            print("{} namespace {}".format(etype, obj.metadata.name))
 
 
 async def watch_pods():
     v1 = client.CoreV1Api()
-    async for event in watch.Watch().stream(v1.list_pod_for_all_namespaces):
-        evt, obj = event['type'], event['object']
-        print("{} pod {} in NS {}".format(evt, obj.metadata.name, obj.metadata.namespace))
+    async with watch.Watch().stream(v1.list_pod_for_all_namespaces) as stream:
+        async for event in stream:
+            evt, obj = event['type'], event['object']
+            print("{} pod {} in NS {}".format(evt, obj.metadata.name, obj.metadata.namespace))
 
 
 def main():
diff --git a/kubernetes_asyncio/watch/watch.py b/kubernetes_asyncio/watch/watch.py
index 1bb131a6..eab69724 100644
--- a/kubernetes_asyncio/watch/watch.py
+++ b/kubernetes_asyncio/watch/watch.py
@@ -50,6 +50,7 @@ class Watch(object):
         self._stop = False
         self._api_client = client.ApiClient()
         self.resource_version = 0
+        self.resp = None
 
     def stop(self):
         self._stop = True
@@ -107,7 +108,11 @@ class Watch(object):
         return self
 
     async def __anext__(self):
-        return await self.next()
+        try:
+            return await self.next()
+        except:
+            self.close()
+            raise
 
     async def next(self):
 
@@ -169,12 +174,23 @@ class Watch(object):
                 if should_stop:
                     watch.stop()
         """
+        self.close()
         self._stop = False
         self.return_type = self.get_return_type(func)
         kwargs['watch'] = True
         kwargs['_preload_content'] = False
 
         self.func = partial(func, *args, **kwargs)
-        self.resp = None
 
         return self
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        self.close()
+    
+    def close(self):
+        if self.resp is not None:
+            self.resp.release()
+            self.resp = None


Test patch:
diff --git a/kubernetes_asyncio/watch/watch_test.py b/kubernetes_asyncio/watch/watch_test.py
index 97459808..5fce0883 100644
--- a/kubernetes_asyncio/watch/watch_test.py
+++ b/kubernetes_asyncio/watch/watch_test.py
@@ -26,6 +26,7 @@ class WatchTest(TestCase):
     async def test_watch_with_decode(self):
         fake_resp = CoroutineMock()
         fake_resp.content.readline = CoroutineMock()
+        fake_resp.release = Mock()
         side_effects = [
             {
                 "type": "ADDED",
@@ -47,22 +48,24 @@ class WatchTest(TestCase):
 
         watch = kubernetes_asyncio.watch.Watch()
         count = 0
-        async for e in watch.stream(fake_api.get_namespaces, resource_version='123'):
-            self.assertEqual("ADDED", e['type'])
-            # make sure decoder worked and we got a model with the right name
-            self.assertEqual("test%d" % count, e['object'].metadata.name)
-            # make sure decoder worked and updated Watch.resource_version
-            self.assertEqual(e['object'].metadata.resource_version, str(count))
-            self.assertEqual(watch.resource_version, str(count))
-
-            # 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()
+        async with watch:
+            async for e in watch.stream(fake_api.get_namespaces, resource_version='123'):
+                self.assertEqual("ADDED", e['type'])
+                # make sure decoder worked and we got a model with the right name
+                self.assertEqual("test%d" % count, e['object'].metadata.name)
+                # make sure decoder worked and updated Watch.resource_version
+                self.assertEqual(e['object'].metadata.resource_version, str(count))
+                self.assertEqual(watch.resource_version, str(count))
+
+                # 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()
 
         fake_api.get_namespaces.assert_called_once_with(
             _preload_content=False, watch=True, resource_version='123')
+        fake_resp.release.assert_called_once_with()
 
     async def test_watch_k8s_empty_response(self):
         """Stop the iterator when the response is empty.
@@ -161,6 +164,7 @@ class WatchTest(TestCase):
     async def test_watch_timeout(self):
         fake_resp = CoroutineMock()
         fake_resp.content.readline = CoroutineMock()
+        fake_resp.release = Mock()
 
         mock_event = {"type": "ADDED",
                       "object": {"metadata": {"name": "test1555",
@@ -177,13 +181,14 @@ class WatchTest(TestCase):
         fake_api.get_namespaces.__doc__ = ':return: V1NamespaceList'
 
         watch = kubernetes_asyncio.watch.Watch()
-        async for e in watch.stream(fake_api.get_namespaces): # noqa
-            pass
+        async with watch.stream(fake_api.get_namespaces) as stream:
+            async for e in stream: # noqa
+                pass
 
         fake_api.get_namespaces.assert_has_calls(
             [call(_preload_content=False, watch=True),
              call(_preload_content=False, watch=True, resource_version='1555')])
-
+        fake_resp.release.assert_called_once_with()
 
 if __name__ == '__main__':
     import asynctest