def reset_environment(user: str, auth: AuthMethod = AuthMethod.ASK):
"""Clears the environment and adds the testing cluster.
:param user: User to connect to the cluster as,
and whose home dir should be cleaned.
:param auth: Authentication method to use.
"""
# pylint: disable=protected-access
saved_state = EnvironmentProvider._state
EnvironmentProvider._state = None
os.environ['IDACT_KEY_LOCATION'] = get_test_key_location(user=user)
os.environ['IDACT_CONFIG_PATH'] = get_test_environment_file(user=user)
cluster = ClusterConfigImpl(host=get_testing_host(),
port=get_testing_port(),
user=user,
auth=auth,
retries=get_default_retries_heavy_load())
cluster.retries[Retry.PORT_INFO] = set_retry(count=0)
EnvironmentProvider(initial_environment=EnvironmentImpl(
config=ClientConfig(clusters={TEST_CLUSTER: cluster})))
set_log_level(DEBUG)
try:
yield
finally:
EnvironmentProvider._state = saved_state
clear_home(user=user)
def test_generate_and_install_key_on_access_node():
with ExitStack() as stack:
user = USER_8
stack.enter_context(clear_environment(user))
stack.enter_context(set_up_key_location(user))
stack.enter_context(disable_pytest_stdin())
with pytest.raises(ValueError): # No key provided.
add_cluster(name=TEST_CLUSTER,
user=user,
host=get_testing_host(),
port=get_testing_port(),
auth=AuthMethod.PUBLIC_KEY,
key=None,
install_key=True)
# Generate RSA key.
add_cluster(name=TEST_CLUSTER,
user=user,
host=get_testing_host(),
port=get_testing_port(),
auth=AuthMethod.PUBLIC_KEY,
key=KeyType.RSA,
install_key=True,
retries={Retry.PORT_INFO: set_retry(count=0)})
check_remote_key_and_node_access(stack=stack, user=user)
def test_generate_and_install_key_no_sshd():
with ExitStack() as stack:
user = USER_14
stack.enter_context(clear_environment(user))
stack.enter_context(set_up_key_location(user))
stack.enter_context(disable_pytest_stdin())
add_cluster(name=TEST_CLUSTER,
user=user,
host=get_testing_host(),
port=get_testing_port(),
auth=AuthMethod.PUBLIC_KEY,
key=KeyType.RSA,
install_key=True,
disable_sshd=True,
retries={Retry.PORT_INFO: set_retry(count=0)})
check_remote_key_and_node_access(stack=stack, user=user)
def test_generate_and_install_missing_key_on_access_node():
with ExitStack() as stack:
user = USER_12
stack.enter_context(clear_environment(user))
stack.enter_context(set_up_key_location(user))
stack.enter_context(disable_pytest_stdin())
missing_key = os.path.join(os.environ['IDACT_KEY_LOCATION'],
'id_rsa_fake')
add_cluster(name=TEST_CLUSTER,
user=user,
host=get_testing_host(),
port=get_testing_port(),
auth=AuthMethod.PUBLIC_KEY,
key=missing_key,
install_key=True,
retries={Retry.PORT_INFO: set_retry(count=0)})
cluster = show_cluster(TEST_CLUSTER)
node = cluster.get_access_node()
generate_missing_key_prompt = (
"Generate new public/private key pair? [Y/n] ")
with set_password(get_test_user_password(user)):
with pytest.raises(IOError): # Will prompt for input.
node.run('whoami')
# Key missing and user chose not to generate one.
with settings(prompts={generate_missing_key_prompt: 'n'}):
with pytest.raises(RuntimeError):
node.run('whoami')
# Key missing, generated and installed.
with settings(prompts={generate_missing_key_prompt: 'y'}):
assert node.run('whoami') == user
check_remote_key_and_node_access(stack=stack, user=user)
def get_default_retries() -> Dict[Retry, RetryConfig]:
"""Returns all default retry values."""
return {
Retry.PORT_INFO: set_retry(count=25, seconds_between=1),
Retry.JUPYTER_JSON: set_retry(count=15, seconds_between=1),
Retry.SCHEDULER_CONNECT: set_retry(count=5, seconds_between=2),
Retry.DASK_NODE_CONNECT: set_retry(count=3, seconds_between=5),
Retry.DEPLOY_DASK_SCHEDULER: set_retry(count=3, seconds_between=5),
Retry.DEPLOY_DASK_WORKER: set_retry(count=3, seconds_between=5),
Retry.GET_SCHEDULER_ADDRESS: set_retry(count=25, seconds_between=1),
Retry.CHECK_WORKER_STARTED: set_retry(count=25, seconds_between=1),
Retry.CANCEL_DEPLOYMENT: set_retry(count=5, seconds_between=1),
Retry.SQUEUE_AFTER_SBATCH: set_retry(count=9, seconds_between=1),
Retry.OPEN_TUNNEL: set_retry(count=3, seconds_between=5),
Retry.VALIDATE_HTTP_TUNNEL: set_retry(count=6, seconds_between=1),
Retry.TUNNEL_TRY_AGAIN_WITH_ANY_PORT: set_retry(count=1,
seconds_between=0)}