add token auth to kernel websocket test

jupyter · Jan 20, 2017 · 2a5f8d2 · 2a5f8d2
1 parent 75a0d71
commit 2a5f8d2
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 11 deletions.
diff --git a/notebook/services/kernels/tests/test_kernels_api.py b/notebook/services/kernels/tests/test_kernels_api.py
@@ -3,9 +3,9 @@
 import json
 import time
 
-import requests
-from tornado.websocket import websocket_connect
+from tornado.httpclient import HTTPRequest
 from tornado.ioloop import IOLoop
+from tornado.websocket import websocket_connect
 
 from jupyter_client.kernelspec import NATIVE_KERNEL_NAME
 
@@ -14,8 +14,10 @@
 
 class KernelAPI(object):
     """Wrapper for kernel REST API requests"""
-    def __init__(self, request):
+    def __init__(self, request, base_url, headers):
         self.request = request
+        self.base_url = base_url
+        self.headers = headers
 
     def _req(self, verb, path, body=None):
         response = self.request(verb,
@@ -51,16 +53,21 @@ def restart(self, id):
 
     def websocket(self, id):
         loop = IOLoop()
-        f = websocket_connect(url_path_join(
-            self.base_url.replace('http', 'ws', 1), 'api/kernels', id, 'channels'),
-            io_loop=loop)
+        req = HTTPRequest(
+            url_path_join(self.base_url.replace('http', 'ws', 1), 'api/kernels', id, 'channels'),
+            headers=self.headers,
+        )
+        f = websocket_connect(req, io_loop=loop)
         return loop.run_sync(lambda : f)
 
 
 class KernelAPITest(NotebookTestBase):
     """Test the kernels web service API"""
     def setUp(self):
-        self.kern_api = KernelAPI(self.request)
+        self.kern_api = KernelAPI(self.request,
+                                  base_url=self.base_url(),
+                                  headers=self.auth_headers(),
+                                  )
 
     def tearDown(self):
         for k in self.kern_api.list().json():

diff --git a/notebook/tests/launchnotebook.py b/notebook/tests/launchnotebook.py
@@ -72,16 +72,22 @@ def wait_until_dead(cls):
             raise TimeoutError("Undead notebook server")
 
     @classmethod
-    def request(self, verb, path, **kwargs):
+    def auth_headers(cls):
+        headers = {}
+        if cls.token:
+            headers['Authorization'] = 'token %s' % cls.token
+        return headers
+
+    @classmethod
+    def request(cls, verb, path, **kwargs):
         """Send a request to my server
         
         with authentication and everything.
         """
         headers = kwargs.setdefault('headers', {})
-        # kwargs.setdefault('allow_redirects', False)
-        headers.setdefault('Authorization', 'token %s' % self.token)
+        headers.update(cls.auth_headers())
         response = requests.request(verb,
-            url_path_join(self.base_url(), path),
+            url_path_join(cls.base_url(), path),
             **kwargs)
         return response