Improvements on ws_client. Now the client can returns an object to interact with websocket server and reach each channel separately

mbohlool · mbohlool · commit d1d2507ca5e4 · 2017-02-17T13:55:33.000-08:00
diff --git a/examples/exec.py b/examples/exec.py
@@ -0,0 +1,94 @@
+import time
+
+from kubernetes import config
+from kubernetes.client import configuration
+from kubernetes.client.apis import core_v1_api
+from kubernetes.client.rest import ApiException
+
+config.load_kube_config()
+configuration.assert_hostname = False
+api = core_v1_api.CoreV1Api()
+name = 'busybox-test'
+
+resp = None
+try:
+    resp = api.read_namespaced_pod(name=name,
+                                   namespace='default')
+except ApiException as e:
+    if e.status != 404:
+        print("Unknown error: %s" % e)
+        exit(1)
+
+if not resp:
+    print("Pod %s does not exits. Creating it..." % name)
+    pod_manifest = {
+        'apiVersion': 'v1',
+        'kind': 'Pod',
+        'metadata': {
+            'name': name
+        },
+        'spec': {
+            'containers': [{
+                'image': 'busybox',
+                'name': 'sleep',
+                "args": [
+                    "/bin/sh",
+                    "-c",
+                    "while true;do date;sleep 5; done"
+                ]
+            }]
+        }
+    }
+    resp = api.create_namespaced_pod(body=pod_manifest,
+                                     namespace='default')
+    while True:
+        resp = api.read_namespaced_pod(name=name,
+                                       namespace='default')
+        if resp.status.phase != 'Pending':
+            break
+        time.sleep(1)
+    print("Done.")
+
+
+# calling exec and wait for response.
+exec_command = [
+    '/bin/sh',
+    '-c',
+    'echo This message goes to stderr >&2; echo This message goes to stdout']
+resp = api.connect_get_namespaced_pod_exec(name, 'default',
+                                           command=exec_command,
+                                           stderr=True, stdin=False,
+                                           stdout=True, tty=False)
+print("Response: " + resp)
+
+# Calling exec interactively.
+exec_command = ['/bin/sh']
+resp = api.connect_get_namespaced_pod_exec(name, 'default',
+                                           command=exec_command,
+                                           stderr=True, stdin=True,
+                                           stdout=True, tty=False,
+
+                                           _preload_content=False)
+commands = [
+    "echo test1",
+    "echo \"This message goes to stderr\" >&2",
+]
+while resp.is_open():
+    resp.update(timeout=1)
+    if resp.peek_stdout():
+        print("STDOUT: %s" % resp.read_stdout())
+    if resp.peek_stderr():
+        print("STDERR: %s" % resp.read_stderr())
+    if commands:
+        c = commands.pop(0)
+        print("Running command... %s\n" % c)
+        resp.write_stdin(c + "\n")
+    else:
+        break
+
+resp.write_stdin("date\n")
+sdate = resp.readline_stdout(timeout=3)
+print("Server date command returns: %s" % sdate)
+resp.write_stdin("whoami\n")
+user = resp.readline_stdout(timeout=3)
+print("Server user is: %s" % user)
diff --git a/kubernetes/client/api_client.py b/kubernetes/client/api_client.py
@@ -351,6 +351,7 @@ def request(self, method, url, query_params=None, headers=None,
                                  url,
                                  query_params=query_params,
                                  _request_timeout=_request_timeout,
+                                 _preload_content=_preload_content,
                                  headers=headers)
 
         if method == "GET":
diff --git a/kubernetes/client/ws_client.py b/kubernetes/client/ws_client.py
@@ -12,33 +12,33 @@
 
 from .rest import ApiException
 
+import select
 import certifi
+import time
 import collections
-import websocket
+from websocket import WebSocket, ABNF, enableTrace
 import six
 import ssl
 from six.moves.urllib.parse import urlencode
 from six.moves.urllib.parse import quote_plus
 
+STDIN_CHANNEL = 0
+STDOUT_CHANNEL = 1
+STDERR_CHANNEL = 2
+
 
 class WSClient:
     def __init__(self, configuration, url, headers):
-        self.messages = []
-        self.errors = []
-        websocket.enableTrace(False)
-        header = None
+        enableTrace(False)
+        header = []
+        self._connected = False
+        self._channels = {}
+        self._all = ""
 
         # We just need to pass the Authorization, ignore all the other
         # http headers we get from the generated code
         if 'Authorization' in headers:
-            header = "Authorization: %s" % headers['Authorization']
-
-        self.ws = websocket.WebSocketApp(url,
-                                         on_message=self.on_message,
-                                         on_error=self.on_error,
-                                         on_close=self.on_close,
-                                         header=[header] if header else None)
-        self.ws.on_open = self.on_open
+            header.append("Authorization: %s" % headers['Authorization'])
 
         if url.startswith('wss://') and configuration.verify_ssl:
             ssl_opts = {
@@ -52,30 +52,118 @@ def __init__(self, configuration, url, headers):
         else:
             ssl_opts = {'cert_reqs': ssl.CERT_NONE}
 
-        self.ws.run_forever(sslopt=ssl_opts)
-
-    def on_message(self, ws, message):
-        if message[0] == '\x01':
-            message = message[1:]
-        if message:
-            if six.PY3 and isinstance(message, six.binary_type):
-                message = message.decode('utf-8')
-            self.messages.append(message)
+        self.sock = WebSocket(sslopt=ssl_opts, skip_utf8_validation=False)
+        self.sock.connect(url, header=header)
+        self._connected = True
 
-    def on_error(self, ws, error):
-        self.errors.append(error)
+    def peek_channel(self, channel, timeout=0):
+        self.update(timeout=timeout)
+        if channel in self._channels:
+            return self._channels[channel]
+        return ""
 
-    def on_close(self, ws):
-        pass
-
-    def on_open(self, ws):
-        pass
+    def read_channel(self, channel, timeout=0):
+        if channel not in self._channels:
+            ret = self.peek_channel(channel, timeout)
+        else:
+            ret = self._channels[channel]
+        if channel in self._channels:
+            del self._channels[channel]
+        return ret
+
+    def readline_channel(self, channel, timeout=None):
+        if timeout is None:
+            timeout = float("inf")
+        start = time.time()
+        while self.is_open() and time.time() - start < timeout:
+            if channel in self._channels:
+                data = self._channels[channel]
+                if "\n" in data:
+                    index = data.find("\n")
+                    ret = data[:index]
+                    data = data[index+1:]
+                    if data:
+                        self._channels[channel] = data
+                    else:
+                        del self._channels[channel]
+                    return ret
+            self.update(timeout=(timeout - time.time() + start))
+
+    def write_channel(self, channel, data):
+        self.sock.send(chr(channel) + data)
+
+    def peek_stdout(self, timeout=0):
+        return self.peek_channel(STDOUT_CHANNEL, timeout=timeout)
+
+    def read_stdout(self, timeout=None):
+        return self.read_channel(STDOUT_CHANNEL, timeout=timeout)
+
+    def readline_stdout(self, timeout=None):
+        return self.readline_channel(STDOUT_CHANNEL, timeout=timeout)
+
+    def peek_stderr(self, timeout=0):
+        return self.peek_channel(STDERR_CHANNEL, timeout=timeout)
+
+    def read_stderr(self, timeout=None):
+        return self.read_channel(STDERR_CHANNEL, timeout=timeout)
+
+    def readline_stderr(self, timeout=None):
+        return self.readline_channel(STDERR_CHANNEL, timeout=timeout)
+
+    def read_all(self):
+        out = self._all
+        self._all = ""
+        self._channels = {}
+        return out
+
+    def is_open(self):
+        return self._connected
+
+    def write_stdin(self, data):
+        self.write_channel(STDIN_CHANNEL, data)
+
+    def update(self, timeout=0):
+        if not self.is_open():
+            return
+        if not self.sock.connected:
+            self._connected = False
+            return
+        r, _, _ = select.select(
+            (self.sock.sock, ), (), (), timeout)
+        if r:
+            op_code, frame = self.sock.recv_data_frame(True)
+            if op_code == ABNF.OPCODE_CLOSE:
+                self._connected = False
+                return
+            elif op_code == ABNF.OPCODE_BINARY or op_code == ABNF.OPCODE_TEXT:
+                data = frame.data
+                if six.PY3:
+                    data = data.decode("utf-8")
+                self._all += data
+                if len(data) > 1:
+                    channel = ord(data[0])
+                    data = data[1:]
+                    if data:
+                        if channel not in self._channels:
+                            self._channels[channel] = data
+                        else:
+                            self._channels[channel] += data
+
+    def run_forever(self, timeout=None):
+        if timeout:
+            start = time.time()
+            while self.is_open() and time.time() - start < timeout:
+                self.update(timeout=(timeout - time.time() + start))
+        else:
+            while self.is_open():
+                self.update(timeout=None)
 
 
 WSResponse = collections.namedtuple('WSResponse', ['data'])
 
 
-def GET(configuration, url, query_params, _request_timeout, headers):
+def GET(configuration, url, query_params, _request_timeout, _preload_content,
+        headers):
     # switch protocols from http to websocket
     url = url.replace('http://', 'ws://')
     url = url.replace('https://', 'wss://')
@@ -105,10 +193,11 @@ def GET(configuration, url, query_params, _request_timeout, headers):
     else:
         url += '&command=' + quote_plus(commands)
 
-    client = WSClient(configuration, url, headers)
-    if client.errors:
-        raise ApiException(
-            status=0,
-            reason='\n'.join([str(error) for error in client.errors])
-        )
-    return WSResponse('%s' % ''.join(client.messages))
+    try:
+        client = WSClient(configuration, url, headers)
+        if not _preload_content:
+            return client
+        client.run_forever(timeout=_request_timeout)
+        return WSResponse('%s' % ''.join(client.read_all()))
+    except (Exception, KeyboardInterrupt, SystemExit) as e:
+        raise ApiException(status=0, reason=str(e))
diff --git a/kubernetes/e2e_test/base.py b/kubernetes/e2e_test/base.py
@@ -42,4 +42,5 @@ def get_e2e_configuration():
     if config.host is None:
         raise unittest.SkipTest('Unable to find a running Kubernetes instance')
     print('Running test against : %s' % config.host)
+    config.assert_hostname = False
     return config
diff --git a/kubernetes/e2e_test/test_client.py b/kubernetes/e2e_test/test_client.py
