def test_environment():
user = USER_2
test_environment_file = get_test_environment_file(user=user)
with ExitStack() as stack:
stack.enter_context(clear_environment(user))
stack.enter_context(set_password(get_test_user_password(user)))
clusters = show_clusters()
assert clusters == {}
cluster = add_cluster(name=TEST_CLUSTER,
user=user,
host='localhost',
port=2222)
clusters = show_clusters()
assert show_cluster(name=TEST_CLUSTER) is cluster
assert len(show_clusters()) == 1
assert clusters[TEST_CLUSTER] == cluster
assert cluster.name == TEST_CLUSTER
try:
save_environment(path=test_environment_file)
with clear_environment(user):
assert show_clusters() == {}
load_environment(path=test_environment_file)
cluster2 = show_cluster(name=TEST_CLUSTER)
assert cluster2 is not cluster
assert cluster2 == cluster
finally:
os.remove(test_environment_file)
def test_environment_missing_and_defaults():
user = USER_25
with ExitStack() as stack:
stack.enter_context(clear_environment(user))
EnvironmentProvider()._environment = None # noqa, pylint: disable=protected-access,line-too-long
assert show_clusters() == {}
assert not os.path.isfile(os.environ['IDACT_CONFIG_PATH'])
with pytest.raises(ValueError):
load_environment()
add_cluster(name=TEST_CLUSTER, user=user, host='localhost')
config = show_cluster(TEST_CLUSTER).config
set_log_level(logging.DEBUG)
check_config_is_default(config=config, user=user)
try:
save_environment()
with open(os.environ['IDACT_CONFIG_PATH'], 'r') as test_file:
contents = test_file.read().splitlines()
pprint(contents)
assert contents == get_default_config_contents(user=user)
finally:
os.remove(os.environ['IDACT_CONFIG_PATH'])
def test_environment_create_modify_save_load():
user = USER_19
test_environment_file = get_test_environment_file(user=user)
test_environment_file2 = get_test_environment_file(user=user) + '_2'
with ExitStack() as stack:
stack.enter_context(clear_environment(user))
assert show_clusters() == {}
add_cluster(name=TEST_CLUSTER, user=user, host='localhost')
config = show_cluster(TEST_CLUSTER).config
check_config_is_default(config=config, user=user)
try:
save_environment(path=test_environment_file)
with open(test_environment_file, 'r') as test_file:
contents = test_file.read().splitlines()
pprint(contents)
assert contents == get_default_config_contents(user=user)
config = show_cluster(TEST_CLUSTER).config
config.host = 'localhost2'
config.port = 2222
config.user = '******'
config.auth = AuthMethod.PUBLIC_KEY
config.key = './fake-key'
config.install_key = False
config.disable_sshd = True
config.setup_actions.jupyter = ['abc']
config.setup_actions.dask = ['abc', 'def']
config.scratch = '$HOME2'
config.use_jupyter_lab = False
set_log_level(logging.INFO)
check_config_is_modified(config=config)
save_environment(path=test_environment_file2)
with open(test_environment_file2, 'r') as test_file:
contents = test_file.read().splitlines()
pprint(contents)
assert contents == get_modified_config_contents()
load_environment(test_environment_file)
config = show_cluster(TEST_CLUSTER).config
check_config_is_default(config=config, user=user)
load_environment(test_environment_file2)
config = show_cluster(TEST_CLUSTER).config
check_config_is_modified(config=config)
finally:
try:
os.remove(test_environment_file)
finally:
os.remove(test_environment_file2)
def run_notebook_app(user: str,
environment_file: str,
args: List[str],
notebook_defaults: dict = None) -> Result:
"""Runs the notebook app. Returns the result from CliRunner."""
with ExitStack() as stack:
stack.enter_context(disable_pytest_stdin())
stack.enter_context(set_up_key_location(user))
stack.enter_context(
reset_environment(user=user, auth=AuthMethod.PUBLIC_KEY))
stack.enter_context(set_password(get_test_user_password(user)))
cluster = show_cluster(name=TEST_CLUSTER)
cluster.get_access_node().connect()
set_log_level(logging.INFO)
if notebook_defaults:
config = cluster.config
assert isinstance(config, ClusterConfigImpl)
config.notebook_defaults = notebook_defaults
save_environment(path=environment_file)
try:
runner = CliRunner()
opened_in_browser = []
def fake_open_in_browser(_):
opened_in_browser.append(True)
saved_open_in_browser = JupyterDeploymentImpl.open_in_browser
JupyterDeploymentImpl.open_in_browser = fake_open_in_browser
try:
os.environ[NOTEBOOK_TEST_RUN_ENVIRONMENT_VARIABLE] = ''
result = runner.invoke(main, args=args)
finally:
JupyterDeploymentImpl.open_in_browser = saved_open_in_browser
del os.environ[NOTEBOOK_TEST_RUN_ENVIRONMENT_VARIABLE]
print("\n\n\nClick output of the notebook app run:")
print(result.output)
yield result
finally:
os.remove(environment_file)
def main(cluster_name: str, environment: Optional[str], save_defaults: bool,
reset_defaults: bool, nodes: Optional[int], cores: Optional[int],
memory_per_node: Optional[str], walltime: Optional[str],
native_arg: List[Tuple[str, str]]) -> int:
"""A console script that executes a Jupyter Notebook instance on
an allocated cluster node, and makes it accessible
in the local browser.
CLUSTER_NAME argument is the cluster name to execute the notebook on.
It must already be present in the config file.
"""
ensure_stdin_has_fileno()
log = None
try:
with ExitStack() as stack:
click.echo("Loading environment.")
load_environment(path=environment)
log = get_logger(__name__)
cluster = show_cluster(name=cluster_name)
config = cluster.config
assert isinstance(config, ClusterConfigImpl)
if reset_defaults:
click.echo("Resetting allocation parameters to defaults.")
config.notebook_defaults = {}
parameters = AppAllocationParameters.deserialize(
serialized=config.notebook_defaults)
override_parameters_if_possible(parameters=parameters,
nodes=nodes,
cores=cores,
memory_per_node=memory_per_node,
walltime=walltime,
native_args=native_arg)
if save_defaults:
click.echo("Saving defaults.")
config.notebook_defaults = parameters.serialize()
save_environment(path=environment)
click.echo(format_allocation_parameters(parameters=parameters))
click.echo("Allocating nodes.")
nodes = cluster.allocate_nodes(
nodes=parameters.nodes,
cores=parameters.cores,
memory_per_node=parameters.memory_per_node,
walltime=parameters.walltime,
native_args=convert_native_args_from_command_line_to_dict(
native_args=parameters.native_args))
stack.enter_context(cancel_on_exit(nodes))
nodes.wait()
notebook = nodes[0].deploy_notebook()
stack.enter_context(cancel_local_on_exit(notebook))
click.echo("Pushing the allocation deployment.")
cluster.push_deployment(nodes)
click.echo("Pushing the notebook deployment.")
cluster.push_deployment(notebook)
click.echo(format_deployments_info(cluster_name=cluster_name))
click.echo("Notebook address: ", nl=False)
click.echo(click.style(notebook.address, fg='red'))
notebook.open_in_browser()
sleep_until_allocation_ends(nodes=nodes)
except: # noqa, pylint: disable=broad-except
if log is not None:
log.error("Exception raised.", exc_info=1)
return 1
raise
return 0