File: test_kernels_api.py

package info (click to toggle)
jupyter-notebook 6.4.13-5
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 13,860 kB
  • sloc: javascript: 20,765; python: 15,658; makefile: 255; sh: 160
file content (285 lines) | stat: -rw-r--r-- 9,675 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
"""Test the kernels service API."""

import json
import time

from requests import HTTPError
from traitlets.config import Config

from tornado.httpclient import HTTPRequest
from tornado.ioloop import IOLoop
from tornado.websocket import websocket_connect
from unittest import SkipTest

from jupyter_client.kernelspec import NATIVE_KERNEL_NAME

from notebook.utils import url_path_join
from notebook.tests.launchnotebook import NotebookTestBase, assert_http_error

try:
    from jupyter_client import AsyncMultiKernelManager
    async_testing_enabled = True
except ImportError:
    async_testing_enabled = False


class KernelAPI:
    """Wrapper for kernel REST API requests"""
    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,
                url_path_join('api/kernels', path), data=body)

        if 400 <= response.status_code < 600:
            try:
                response.reason = response.json()['message']
            except:
                pass
        response.raise_for_status()

        return response

    def list(self):
        return self._req('GET', '')

    def get(self, id):
        return self._req('GET', id)

    def start(self, name=NATIVE_KERNEL_NAME):
        body = json.dumps({'name': name})
        return self._req('POST', '', body)

    def shutdown(self, id):
        return self._req('DELETE', id)

    def interrupt(self, id):
        return self._req('POST', url_path_join(id, 'interrupt'))

    def restart(self, id):
        return self._req('POST', url_path_join(id, 'restart'))

    def websocket(self, id):
        loop = IOLoop()
        loop.make_current()
        req = HTTPRequest(
            url_path_join(self.base_url.replace('http', 'ws', 1), 'api/kernels', id, 'channels'),
            headers=self.headers,
        )
        f = websocket_connect(req)
        return loop.run_sync(lambda : f)


class KernelAPITest(NotebookTestBase):
    """Test the kernels web service API"""
    def setUp(self):
        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():
            self.kern_api.shutdown(k['id'])

    def test_no_kernels(self):
        """Make sure there are no kernels running at the start"""
        kernels = self.kern_api.list().json()
        self.assertEqual(kernels, [])

    def test_default_kernel(self):
        # POST request
        r = self.kern_api._req('POST', '')
        kern1 = r.json()
        self.assertEqual(r.headers['location'], url_path_join(self.url_prefix, 'api/kernels', kern1['id']))
        self.assertEqual(r.status_code, 201)
        self.assertIsInstance(kern1, dict)

        report_uri = url_path_join(self.url_prefix, 'api/security/csp-report')
        expected_csp = '; '.join([
            "frame-ancestors 'self'",
            'report-uri ' + report_uri,
            "default-src 'none'"
        ])
        self.assertEqual(r.headers['Content-Security-Policy'], expected_csp)

    def test_main_kernel_handler(self):
        # POST request
        r = self.kern_api.start()
        kern1 = r.json()
        self.assertEqual(r.headers['location'], url_path_join(self.url_prefix, 'api/kernels', kern1['id']))
        self.assertEqual(r.status_code, 201)
        self.assertIsInstance(kern1, dict)

        report_uri = url_path_join(self.url_prefix, 'api/security/csp-report')
        expected_csp = '; '.join([
            "frame-ancestors 'self'",
            'report-uri ' + report_uri,
            "default-src 'none'"
        ])
        self.assertEqual(r.headers['Content-Security-Policy'], expected_csp)

        # GET request
        r = self.kern_api.list()
        self.assertEqual(r.status_code, 200)
        assert isinstance(r.json(), list)
        self.assertEqual(r.json()[0]['id'], kern1['id'])
        self.assertEqual(r.json()[0]['name'], kern1['name'])

        # create another kernel and check that they both are added to the
        # list of kernels from a GET request
        kern2 = self.kern_api.start().json()
        assert isinstance(kern2, dict)
        r = self.kern_api.list()
        kernels = r.json()
        self.assertEqual(r.status_code, 200)
        assert isinstance(kernels, list)
        self.assertEqual(len(kernels), 2)

        # Interrupt a kernel
        r = self.kern_api.interrupt(kern2['id'])
        self.assertEqual(r.status_code, 204)

        # Restart a kernel
        r = self.kern_api.restart(kern2['id'])
        rekern = r.json()
        self.assertEqual(rekern['id'], kern2['id'])
        self.assertEqual(rekern['name'], kern2['name'])

    def test_kernel_handler(self):
        # GET kernel with given id
        kid = self.kern_api.start().json()['id']
        r = self.kern_api.get(kid)
        kern1 = r.json()
        self.assertEqual(r.status_code, 200)
        assert isinstance(kern1, dict)
        self.assertIn('id', kern1)
        self.assertEqual(kern1['id'], kid)

        # Request a bad kernel id and check that a JSON
        # message is returned!
        bad_id = '111-111-111-111-111'
        with assert_http_error(404, 'Kernel does not exist: ' + bad_id):
            self.kern_api.get(bad_id)

        # DELETE kernel with id
        r = self.kern_api.shutdown(kid)
        self.assertEqual(r.status_code, 204)
        kernels = self.kern_api.list().json()
        self.assertEqual(kernels, [])

        # Request to delete a non-existent kernel id
        bad_id = '111-111-111-111-111'
        with assert_http_error(404, 'Kernel does not exist: ' + bad_id):
            self.kern_api.shutdown(bad_id)

    def test_connections(self):
        kid = self.kern_api.start().json()['id']
        model = self.kern_api.get(kid).json()
        self.assertEqual(model['connections'], 0)

        ws = self.kern_api.websocket(kid)
        model = self.kern_api.get(kid).json()
        self.assertEqual(model['connections'], 1)
        ws.close()
        # give it some time to close on the other side:
        for i in range(10):
            model = self.kern_api.get(kid).json()
            if model['connections'] > 0:
                time.sleep(0.1)
            else:
                break
        model = self.kern_api.get(kid).json()
        self.assertEqual(model['connections'], 0)


class AsyncKernelAPITest(KernelAPITest):
    """Test the kernels web service API using the AsyncMappingKernelManager"""

    @classmethod
    def setup_class(cls):
        if not async_testing_enabled:  # Can be removed once jupyter_client >= 6.1 is required.
            raise SkipTest("AsyncKernelAPITest tests skipped due to down-level jupyter_client!")
        super().setup_class()

    @classmethod
    def get_argv(cls):
        argv = super().get_argv()

        # before we extend the argv with the class, ensure that appropriate jupyter_client is available.
        # if not available, don't set kernel_manager_class, resulting in the repeat of sync-based tests.
        if async_testing_enabled:
            argv.extend(['--NotebookApp.kernel_manager_class='
                        'notebook.services.kernels.kernelmanager.AsyncMappingKernelManager'])
        return argv


class KernelFilterTest(NotebookTestBase):

    # A special install of NotebookTestBase where only `kernel_info_request`
    # messages are allowed.
    config = Config({
        'NotebookApp': {
            'MappingKernelManager': {
                'allowed_message_types': ['kernel_info_request']
            }
        }
    })

    # Sanity check verifying that the configurable was properly set.
    def test_config(self):
        self.assertEqual(self.notebook.kernel_manager.allowed_message_types, ['kernel_info_request'])


CULL_TIMEOUT = 5
CULL_INTERVAL = 1


class KernelCullingTest(NotebookTestBase):
    """Test kernel culling """

    @classmethod
    def get_argv(cls):
        argv = super().get_argv()

        # Enable culling with 5s timeout and 1s intervals
        argv.extend([f'--MappingKernelManager.cull_idle_timeout={CULL_TIMEOUT}',
                     f'--MappingKernelManager.cull_interval={CULL_INTERVAL}',
                     '--MappingKernelManager.cull_connected=False'])
        return argv

    def setUp(self):
        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():
            self.kern_api.shutdown(k['id'])

    def test_culling(self):
        kid = self.kern_api.start().json()['id']
        ws = self.kern_api.websocket(kid)
        model = self.kern_api.get(kid).json()
        self.assertEqual(model['connections'], 1)
        assert not self.get_cull_status(kid)  # connected, should not be culled
        ws.close()
        assert self.get_cull_status(kid)  # not connected, should be culled

    def get_cull_status(self, kid):
        frequency = 0.5
        culled = False
        for _ in range(int((CULL_TIMEOUT + CULL_INTERVAL)/frequency)):  # Timeout + Interval will ensure cull
            try:
                self.kern_api.get(kid)
            except HTTPError as e:
                assert e.response.status_code == 404
                culled = True
                break
            else:
                time.sleep(frequency)
        return culled