async def test_basic_auth_password():
config = Config()
config.DaskGateway.authenticator_class = (
"dask_gateway_server.auth.SimpleAuthenticator")
config.SimpleAuthenticator.password = "******"
async with temp_gateway(config=config) as g:
auth = BasicAuth()
async with g.gateway_client(auth=auth) as gateway:
with pytest.raises(Exception):
await gateway.list_clusters()
auth.password = "******"
await gateway.list_clusters()
def test_get_auth():
# Pass through existing auth objects
auth = BasicAuth()
assert get_auth(auth) is auth
# Auth by keyword name
auth = get_auth("basic")
assert isinstance(auth, BasicAuth)
auth = get_auth("kerberos")
assert isinstance(auth, KerberosAuth)
# Auth from config
config = {"gateway": {"auth": {"type": "basic", "kwargs": {}}}}
with dask.config.set(config):
auth = get_auth()
assert isinstance(auth, BasicAuth)
# Auth from config with import path
config = {
"gateway": {
"auth": {
"type": "dask_gateway.auth.BasicAuth",
"kwargs": {}
}
}
}
with dask.config.set(config):
auth = get_auth()
assert isinstance(auth, BasicAuth)
# Auth from config with kwargs
config = {
"gateway": {
"auth": {
"type": "basic",
"kwargs": {
"username": "******"
}
}
}
}
with dask.config.set(config):
auth = get_auth()
assert isinstance(auth, BasicAuth)
assert auth.username == "bruce"
# Errors
with pytest.raises(TypeError):
get_auth(1)
with pytest.raises(TypeError):
get_auth(lambda: 1)
with pytest.raises(ImportError):
get_auth("dask_gateway.auth.Foo")
with pytest.raises(ImportError):
get_auth("not_a_real_module_name_foo_barrr")
async def test_yarn_backend():
c = Config()
c.YarnClusterConfig.scheduler_cmd = "/opt/miniconda/bin/dask-gateway-scheduler"
c.YarnClusterConfig.worker_cmd = "/opt/miniconda/bin/dask-gateway-worker"
c.YarnClusterConfig.scheduler_memory = "512M"
c.YarnClusterConfig.worker_memory = "512M"
c.YarnClusterConfig.scheduler_cores = 1
c.YarnClusterConfig.worker_cores = 1
c.YarnBackend.keytab = "/home/dask/dask.keytab"
c.YarnBackend.principal = "dask"
c.DaskGateway.backend_class = YarnTestingBackend
async with temp_gateway(config=c) as g:
auth = BasicAuth(username="******")
async with g.gateway_client(auth=auth) as gateway:
async with gateway.new_cluster() as cluster:
db_cluster = g.gateway.backend.db.get_cluster(cluster.name)
res = await g.gateway.backend.do_check_clusters([db_cluster])
assert res == [True]
await cluster.scale(2)
await wait_for_workers(cluster, exact=2)
await cluster.scale(1)
await wait_for_workers(cluster, exact=1)
db_workers = list(db_cluster.workers.values())
async def test():
res = await g.gateway.backend.do_check_workers(db_workers)
assert sum(res) == 1
await with_retries(test, 30, 0.25)
async with cluster.get_client(set_as_default=False) as client:
res = await client.submit(lambda x: x + 1, 1)
assert res == 2
await cluster.scale(0)
await wait_for_workers(cluster, exact=0)
async def test():
res = await g.gateway.backend.do_check_workers(db_workers)
assert sum(res) == 0
await with_retries(test, 30, 0.25)
# No-op for shutdown of already shutdown worker
async def test():
res = await g.gateway.backend.do_check_clusters([db_cluster])
assert res == [False]
await with_retries(test, 30, 0.25)
async def test_basic_auth_password(tmpdir):
config = Config()
config.DaskGateway.temp_dir = str(tmpdir.join("dask-gateway"))
config.DaskGateway.authenticator_class = (
"dask_gateway_server.auth.DummyAuthenticator")
config.DummyAuthenticator.password = "******"
async with temp_gateway(config=config) as gateway_proc:
auth = BasicAuth()
async with Gateway(address=gateway_proc.public_url,
asynchronous=True,
auth=auth) as gateway:
with pytest.raises(Exception):
await gateway.list_clusters()
auth.password = "******"
await gateway.list_clusters()
async def test_client_fetch_timeout():
async with slow_server() as server:
gateway = Gateway(server.address, auth=BasicAuth("alice"))
with pytest.raises(TimeoutError):
await gateway._fetch(HTTPRequest(url=server.address, request_timeout=1))