Repository: tomplus/kubernetes_asyncio (language: python)

Issue:
Allow for configurable heartbeat in WsApiClient request method
When running long commands without intermediate output through the WsApiClient, the connection is closed after timeout.
To correct this error, a custom heartbeat parameter can be passed  to the `ws_connect` method of the `rest_client.pool_manager` attribute in the `request` method of the `WsApiClient` object (lines 80 and 94 of the `stream/ws_client.py` file). Something like:

```python
return await self.rest_client.pool_manager.ws_connect(url, headers=headers, heartbeat=heartbeat)
```

It would be useful to have a way to percolate a custom `Optional[float]` heartbeat value (default `None`) through the `WsApiClient` constructor, such that the user can configure it whenever needed.


Discussion / hints:
It sounds good to me. Are you going to add a pull-request with it?
Hi @tomplus. Yes I have no problems in creating a PR.
I've read the Contribution guidelines but it is not clear to me how to proceed. The document says that I should sign a Contributor License Agreement as stated "above" but there is no "above" in the document. Plus, I don't know if I only have to submit the change or if I should also add some kind of no regression test somewhere.
BTW I will start making the main change and then I'll wait for your answer before submitting the PR.


Fix patch:
diff --git a/kubernetes_asyncio/stream/ws_client.py b/kubernetes_asyncio/stream/ws_client.py
index ddd84c70b..5337d5c00 100644
--- a/kubernetes_asyncio/stream/ws_client.py
+++ b/kubernetes_asyncio/stream/ws_client.py
@@ -49,6 +49,11 @@ def getheader(self, name, default=None):
 
 class WsApiClient(ApiClient):
 
+    def __init__(self, configuration=None, header_name=None, header_value=None,
+                 cookie=None, pool_threads=1, heartbeat=None):
+        super().__init__(configuration, header_name, header_value, cookie, pool_threads)
+        self.heartbeat = heartbeat
+
     async def request(self, method, url, query_params=None, headers=None,
                       post_params=None, body=None, _preload_content=True,
                       _request_timeout=None):
@@ -77,7 +82,7 @@ async def request(self, method, url, query_params=None, headers=None,
         if _preload_content:
 
             resp_all = ''
-            async with self.rest_client.pool_manager.ws_connect(url, headers=headers) as ws:
+            async with self.rest_client.pool_manager.ws_connect(url, headers=headers, heartbeat=self.heartbeat) as ws:
                 async for msg in ws:
                     msg = msg.data.decode('utf-8')
                     if len(msg) > 1:
@@ -91,4 +96,4 @@ async def request(self, method, url, query_params=None, headers=None,
 
         else:
 
-            return await self.rest_client.pool_manager.ws_connect(url, headers=headers)
+            return await self.rest_client.pool_manager.ws_connect(url, headers=headers, heartbeat=self.heartbeat)


Test patch:
diff --git a/kubernetes_asyncio/stream/ws_client_test.py b/kubernetes_asyncio/stream/ws_client_test.py
index 7bd37094e..5d6f22d77 100644
--- a/kubernetes_asyncio/stream/ws_client_test.py
+++ b/kubernetes_asyncio/stream/ws_client_test.py
@@ -19,6 +19,26 @@
 from kubernetes_asyncio.stream.ws_client import WsResponse, get_websocket_url
 
 
+class WsMock:
+    def __init__(self):
+        self.iter = 0
+
+    def __aiter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return self
+
+    async def __aenter__(self):
+        return self
+
+    async def __anext__(self):
+        self.iter += 1
+        if self.iter > 5:
+            raise StopAsyncIteration
+        return WsResponse(200, (chr(1) + 'mock').encode('utf-8'))
+
+
 class WSClientTest(TestCase):
 
     def test_websocket_client(self):
@@ -35,36 +55,43 @@ def test_websocket_client(self):
             self.assertEqual(get_websocket_url(url), ws_url)
 
     async def test_exec_ws(self):
+        mock = CoroutineMock()
+        mock.RESTClientObject.return_value.pool_manager = mock
+        mock.ws_connect.return_value = WsMock()
+        with patch('kubernetes_asyncio.client.api_client.rest', mock):
 
-        class WsMock:
-            def __init__(self):
-                self.iter = 0
-
-            def __aiter__(self):
-                return self
-
-            async def __aexit__(self, exc_type, exc, tb):
-                return self
-
-            async def __aenter__(self):
-                return self
+            api_client = WsApiClient()
+            api_client.configuration.host = 'https://localhost'
+            ws = client.CoreV1Api(api_client=api_client)
+            resp = ws.connect_get_namespaced_pod_exec('pod', 'namespace',
+                                                      command="mock-command",
+                                                      stderr=True, stdin=False,
+                                                      stdout=True, tty=False)
 
-            async def __anext__(self):
-                self.iter += 1
-                if self.iter > 5:
-                    raise StopAsyncIteration
-                return WsResponse(200, (chr(1) + 'mock').encode('utf-8'))
+            ret = await resp
+            self.assertEqual(ret, 'mock' * 5)
+            mock.ws_connect.assert_called_once_with(
+                'wss://localhost/api/v1/namespaces/namespace/pods/pod/exec?'
+                'command=mock-command&stderr=True&stdin=False&stdout=True&tty=False',
+                headers={
+                    'sec-websocket-protocol': 'v4.channel.k8s.io',
+                    'Accept': '*/*',
+                    'User-Agent': api_client.user_agent
+                },
+                heartbeat=None
+            )
 
+    async def test_exec_ws_with_heartbeat(self):
         mock = CoroutineMock()
         mock.RESTClientObject.return_value.pool_manager = mock
         mock.ws_connect.return_value = WsMock()
         with patch('kubernetes_asyncio.client.api_client.rest', mock):
 
-            api_client = WsApiClient()
+            api_client = WsApiClient(heartbeat=30)
             api_client.configuration.host = 'https://localhost'
             ws = client.CoreV1Api(api_client=api_client)
             resp = ws.connect_get_namespaced_pod_exec('pod', 'namespace',
-                                                      command="mock-command",
+                                                      command='mock-command',
                                                       stderr=True, stdin=False,
                                                       stdout=True, tty=False)
 
@@ -77,7 +104,8 @@ async def __anext__(self):
                     'sec-websocket-protocol': 'v4.channel.k8s.io',
                     'Accept': '*/*',
                     'User-Agent': api_client.user_agent
-                }
+                },
+                heartbeat=30
             )