mirror of
https://github.com/jupyterhub/jupyterhub.git
synced 2025-10-13 04:53:01 +00:00
512bbae5cb
- test that .will_resume preserves tokens (worked, but wasn't tested) If a Spawner reuses a token, validate it in the db: - verify that it's in the db - if it doesn't map onto the right user, revoke the token - if it's not in the db, insert it as a user-provided token The most likely case is prior unclean shutdown of something like DockerSpawner, where a spawn failed and thus the token was revoked, but the container was in fact created.
363 lines
9.9 KiB
Python
363 lines
9.9 KiB
Python
"""Tests for process spawning"""
|
|
|
|
# Copyright (c) Jupyter Development Team.
|
|
# Distributed under the terms of the Modified BSD License.
|
|
|
|
import logging
|
|
import os
|
|
import signal
|
|
from subprocess import Popen
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
from tornado import gen
|
|
|
|
from ..objects import Hub, Server
|
|
from .. import orm
|
|
from .. import spawner as spawnermod
|
|
from ..spawner import LocalProcessSpawner, Spawner
|
|
from ..user import User
|
|
from ..utils import new_token
|
|
from .mocking import MockSpawner
|
|
from .test_api import add_user
|
|
from .utils import async_requests
|
|
|
|
_echo_sleep = """
|
|
import sys, time
|
|
print(sys.argv)
|
|
time.sleep(30)
|
|
"""
|
|
|
|
_uninterruptible = """
|
|
import time
|
|
while True:
|
|
try:
|
|
time.sleep(10)
|
|
except KeyboardInterrupt:
|
|
print("interrupted")
|
|
"""
|
|
|
|
|
|
def setup():
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
|
|
|
|
def new_spawner(db, **kwargs):
|
|
user = kwargs.setdefault('user', User(db.query(orm.User).first(), {}))
|
|
kwargs.setdefault('cmd', [sys.executable, '-c', _echo_sleep])
|
|
kwargs.setdefault('hub', Hub())
|
|
kwargs.setdefault('notebook_dir', os.getcwd())
|
|
kwargs.setdefault('default_url', '/user/{username}/lab')
|
|
kwargs.setdefault('oauth_client_id', 'mock-client-id')
|
|
kwargs.setdefault('INTERRUPT_TIMEOUT', 1)
|
|
kwargs.setdefault('TERM_TIMEOUT', 1)
|
|
kwargs.setdefault('KILL_TIMEOUT', 1)
|
|
kwargs.setdefault('poll_interval', 1)
|
|
return user._new_spawner('', spawner_class=LocalProcessSpawner, **kwargs)
|
|
|
|
|
|
@pytest.mark.gen_test
|
|
def test_spawner(db, request):
|
|
spawner = new_spawner(db)
|
|
ip, port = yield spawner.start()
|
|
assert ip == '127.0.0.1'
|
|
assert isinstance(port, int)
|
|
assert port > 0
|
|
db.commit()
|
|
|
|
# wait for the process to get to the while True: loop
|
|
time.sleep(1)
|
|
|
|
status = yield spawner.poll()
|
|
assert status is None
|
|
yield spawner.stop()
|
|
status = yield spawner.poll()
|
|
assert status == 1
|
|
|
|
|
|
@gen.coroutine
|
|
def wait_for_spawner(spawner, timeout=10):
|
|
"""Wait for an http server to show up
|
|
|
|
polling at shorter intervals for early termination
|
|
"""
|
|
deadline = time.monotonic() + timeout
|
|
def wait():
|
|
return spawner.server.wait_up(timeout=1, http=True)
|
|
while time.monotonic() < deadline:
|
|
status = yield spawner.poll()
|
|
assert status is None
|
|
try:
|
|
yield wait()
|
|
except TimeoutError:
|
|
continue
|
|
else:
|
|
break
|
|
yield wait()
|
|
|
|
|
|
@pytest.mark.gen_test
|
|
def test_single_user_spawner(app, request):
|
|
user = next(iter(app.users.values()), None)
|
|
spawner = user.spawner
|
|
spawner.cmd = ['jupyterhub-singleuser']
|
|
yield user.spawn()
|
|
assert spawner.server.ip == '127.0.0.1'
|
|
assert spawner.server.port > 0
|
|
yield wait_for_spawner(spawner)
|
|
status = yield spawner.poll()
|
|
assert status is None
|
|
yield spawner.stop()
|
|
status = yield spawner.poll()
|
|
assert status == 0
|
|
|
|
|
|
@pytest.mark.gen_test
|
|
def test_stop_spawner_sigint_fails(db):
|
|
spawner = new_spawner(db, cmd=[sys.executable, '-c', _uninterruptible])
|
|
yield spawner.start()
|
|
|
|
# wait for the process to get to the while True: loop
|
|
yield gen.sleep(1)
|
|
|
|
status = yield spawner.poll()
|
|
assert status is None
|
|
|
|
yield spawner.stop()
|
|
status = yield spawner.poll()
|
|
assert status == -signal.SIGTERM
|
|
|
|
|
|
@pytest.mark.gen_test
|
|
def test_stop_spawner_stop_now(db):
|
|
spawner = new_spawner(db)
|
|
yield spawner.start()
|
|
|
|
# wait for the process to get to the while True: loop
|
|
yield gen.sleep(1)
|
|
|
|
status = yield spawner.poll()
|
|
assert status is None
|
|
|
|
yield spawner.stop(now=True)
|
|
status = yield spawner.poll()
|
|
assert status == -signal.SIGTERM
|
|
|
|
|
|
@pytest.mark.gen_test
|
|
def test_spawner_poll(db):
|
|
first_spawner = new_spawner(db)
|
|
user = first_spawner.user
|
|
yield first_spawner.start()
|
|
proc = first_spawner.proc
|
|
status = yield first_spawner.poll()
|
|
assert status is None
|
|
if user.state is None:
|
|
user.state = {}
|
|
first_spawner.orm_spawner.state = first_spawner.get_state()
|
|
assert 'pid' in first_spawner.orm_spawner.state
|
|
|
|
# create a new Spawner, loading from state of previous
|
|
spawner = new_spawner(db, user=first_spawner.user)
|
|
spawner.start_polling()
|
|
|
|
# wait for the process to get to the while True: loop
|
|
yield gen.sleep(1)
|
|
status = yield spawner.poll()
|
|
assert status is None
|
|
|
|
# kill the process
|
|
proc.terminate()
|
|
for i in range(10):
|
|
if proc.poll() is None:
|
|
yield gen.sleep(1)
|
|
else:
|
|
break
|
|
assert proc.poll() is not None
|
|
|
|
yield gen.sleep(2)
|
|
status = yield spawner.poll()
|
|
assert status is not None
|
|
|
|
|
|
def test_setcwd():
|
|
cwd = os.getcwd()
|
|
with tempfile.TemporaryDirectory() as td:
|
|
td = os.path.realpath(os.path.abspath(td))
|
|
spawnermod._try_setcwd(td)
|
|
assert os.path.samefile(os.getcwd(), td)
|
|
os.chdir(cwd)
|
|
chdir = os.chdir
|
|
temp_root = os.path.realpath(os.path.abspath(tempfile.gettempdir()))
|
|
def raiser(path):
|
|
path = os.path.realpath(os.path.abspath(path))
|
|
if not path.startswith(temp_root):
|
|
raise OSError(path)
|
|
chdir(path)
|
|
with mock.patch('os.chdir', raiser):
|
|
spawnermod._try_setcwd(cwd)
|
|
assert os.getcwd().startswith(temp_root)
|
|
os.chdir(cwd)
|
|
|
|
|
|
def test_string_formatting(db):
|
|
s = new_spawner(db, notebook_dir='user/%U/', default_url='/base/{username}')
|
|
name = s.user.name
|
|
assert s.notebook_dir == 'user/{username}/'
|
|
assert s.default_url == '/base/{username}'
|
|
assert s.format_string(s.notebook_dir) == 'user/%s/' % name
|
|
assert s.format_string(s.default_url) == '/base/%s' % name
|
|
|
|
|
|
@pytest.mark.gen_test
|
|
def test_popen_kwargs(db):
|
|
mock_proc = mock.Mock(spec=Popen)
|
|
def mock_popen(*args, **kwargs):
|
|
mock_proc.args = args
|
|
mock_proc.kwargs = kwargs
|
|
mock_proc.pid = 5
|
|
return mock_proc
|
|
|
|
s = new_spawner(db, popen_kwargs={'shell': True}, cmd='jupyterhub-singleuser')
|
|
with mock.patch.object(spawnermod, 'Popen', mock_popen):
|
|
yield s.start()
|
|
|
|
assert mock_proc.kwargs['shell'] == True
|
|
assert mock_proc.args[0][:1] == (['jupyterhub-singleuser'])
|
|
|
|
|
|
@pytest.mark.gen_test
|
|
def test_shell_cmd(db, tmpdir, request):
|
|
f = tmpdir.join('bashrc')
|
|
f.write('export TESTVAR=foo\n')
|
|
s = new_spawner(db,
|
|
cmd=[sys.executable, '-m', 'jupyterhub.tests.mocksu'],
|
|
shell_cmd=['bash', '--rcfile', str(f), '-i', '-c'],
|
|
)
|
|
server = orm.Server()
|
|
db.add(server)
|
|
db.commit()
|
|
s.server = Server.from_orm(server)
|
|
db.commit()
|
|
(ip, port) = yield s.start()
|
|
request.addfinalizer(s.stop)
|
|
s.server.ip = ip
|
|
s.server.port = port
|
|
db.commit()
|
|
yield wait_for_spawner(s)
|
|
r = yield async_requests.get('http://%s:%i/env' % (ip, port))
|
|
r.raise_for_status()
|
|
env = r.json()
|
|
assert env['TESTVAR'] == 'foo'
|
|
|
|
|
|
def test_inherit_overwrite():
|
|
"""On 3.6+ we check things are overwritten at import time
|
|
"""
|
|
if sys.version_info >= (3,6):
|
|
with pytest.raises(NotImplementedError):
|
|
class S(Spawner):
|
|
pass
|
|
|
|
|
|
def test_inherit_ok():
|
|
class S(Spawner):
|
|
def start():
|
|
pass
|
|
|
|
def stop():
|
|
pass
|
|
|
|
def poll():
|
|
pass
|
|
|
|
|
|
class MockAPITokenSpawner(MockSpawner):
|
|
"""Mock Spawner that sets its own API token"""
|
|
use_this_api_token = None
|
|
will_resume = True
|
|
def start(self):
|
|
if self.use_this_api_token:
|
|
self.api_token = self.use_this_api_token
|
|
else:
|
|
self.use_this_api_token = self.api_token
|
|
return super().start()
|
|
|
|
|
|
@pytest.mark.gen_test
|
|
def test_spawner_reuse_api_token(db, app):
|
|
user = add_user(app.db, app, name='snoopy')
|
|
if user.running:
|
|
yield user.stop()
|
|
assert user.api_tokens == []
|
|
spawner = user._new_spawner(name='', spawner_class=MockAPITokenSpawner)
|
|
user.spawners[''] = spawner
|
|
assert user.spawner is spawner
|
|
# first start: gets a new API token
|
|
yield user.spawn()
|
|
api_token = spawner.api_token
|
|
found = orm.APIToken.find(app.db, api_token)
|
|
assert found
|
|
assert found.user.name == user.name
|
|
assert user.api_tokens == [found]
|
|
yield user.stop()
|
|
yield user.spawn()
|
|
# verify re-use of API token
|
|
assert spawner.api_token == api_token
|
|
# verify that a new token was not created
|
|
assert user.api_tokens == [found]
|
|
|
|
|
|
@pytest.mark.gen_test
|
|
def test_spawner_insert_api_token(db, app):
|
|
user = add_user(app.db, app, name='tonkee')
|
|
if user.running:
|
|
yield user.stop()
|
|
assert user.api_tokens == []
|
|
|
|
spawner = user._new_spawner(name='', spawner_class=MockAPITokenSpawner)
|
|
user.spawners[''] = spawner
|
|
assert user.spawner is spawner
|
|
|
|
api_token = new_token()
|
|
assert not orm.APIToken.find(app.db, api_token)
|
|
user.spawner.use_this_api_token = api_token
|
|
# Spawner's provided API token is inserted into the db
|
|
# This shouldn't happen unless there are bugs elsewhere (in the Spawner),
|
|
# but handle it anyway.
|
|
yield user.spawn()
|
|
assert spawner.api_token == api_token
|
|
found = orm.APIToken.find(app.db, api_token)
|
|
assert found
|
|
assert found.user.name == user.name
|
|
assert user.api_tokens == [found]
|
|
yield user.stop()
|
|
|
|
|
|
@pytest.mark.gen_test
|
|
def test_spawner_bad_api_token(db, app):
|
|
"""Tokens are revoked when a Spawner gets another user's token"""
|
|
user = add_user(app.db, app, name='hoa')
|
|
if user.running:
|
|
yield user.stop()
|
|
|
|
spawner = user._new_spawner(name='', spawner_class=MockAPITokenSpawner)
|
|
user.spawners[''] = spawner
|
|
assert user.spawner is spawner
|
|
|
|
other_user = add_user(app.db, app, name='greystone')
|
|
assert user.api_tokens == []
|
|
assert other_user.api_tokens == []
|
|
other_token = other_user.new_api_token()
|
|
spawner.use_this_api_token = other_token
|
|
|
|
assert len(other_user.api_tokens) == 1
|
|
|
|
with pytest.raises(ValueError):
|
|
yield user.spawn()
|
|
assert orm.APIToken.find(app.db, other_token) is None
|
|
assert other_user.api_tokens == []
|