Configurable heartbeat value for WsApiClient

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)

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
             )