diff --git a/notebook/services/kernels/tests/test_kernels_api.py b/notebook/services/kernels/tests/test_kernels_api.py index 3b97ec7303..5358a20947 100644 --- 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 index 9ffe34989b..146f633144 100644 --- 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