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)}