def test_environment_provider():
"""Creating and resetting the environment."""
cluster1 = ClusterConfigImpl(host='host1',
port=22,
user='******',
auth=AuthMethod.ASK)
config = ClientConfig(clusters={'cluster1': cluster1})
environment = EnvironmentImpl(config=config)
EnvironmentProvider._state = None # pylint: disable=protected-access
environment_provider = EnvironmentProvider(initial_environment=environment)
assert EnvironmentProvider._state == environment_provider.__dict__ # noqa, pylint: disable=protected-access,line-too-long
assert environment_provider.environment is environment
cluster2 = ClusterConfigImpl(host='host2',
port=22,
user='******',
auth=AuthMethod.ASK)
config2 = ClientConfig(clusters={'cluster2': cluster2})
environment2 = EnvironmentImpl(config=config2)
environment_provider2 = EnvironmentProvider(
initial_environment=environment2)
assert environment_provider2.environment is environment
environment_provider2.environment = environment2
assert environment_provider2.environment is environment2
assert environment_provider.environment is environment2
EnvironmentProvider._state = None # pylint: disable=protected-access
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 __init__(self, config: Optional[ClientConfig] = None):
self._config = ClientConfig(clusters={}) if not config else config
self._clusters = {name: ClusterImpl(name=name,
config=cluster_config)
for name, cluster_config
in self._config.clusters.items()}
self.set_log_level(level=self._config.log_level)
def deserialize_client_config_from_json(data: dict) -> ClientConfig:
"""Deserializes :class:`.ClientConfig` from json.
:param data: json to deserialize.
"""
log = get_logger(__name__)
log.debug("Loaded config: %s", data)
if use_defaults_in_missing_fields(data=data):
log.debug("Filled missing fields with None: %s", data)
clusters = {
name: ClusterConfigImpl(host=value['host'],
port=value['port'],
user=value['user'],
auth=AuthMethod[value['auth']],
key=value['key'],
install_key=value['installKey'],
disable_sshd=value['disableSshd'],
setup_actions=SetupActionsConfigImpl(
jupyter=value['setupActions']['jupyter'],
dask=value['setupActions']['dask']),
scratch=value['scratch'],
notebook_defaults=value['notebookDefaults'],
retries=provide_defaults_for_retries(
deserialize_retries(value['retries'])),
use_jupyter_lab=value['useJupyterLab'])
for name, value in data['clusters'].items()
}
return ClientConfig(clusters=clusters, log_level=data['logLevel'])
def test_client_config_create():
cluster_cluster_config = VALID_CLIENT_CLUSTER_CONFIG
print(cluster_cluster_config.__dict__)
clusters = {'cluster1': cluster_cluster_config}
client_config = ClientConfig(clusters=clusters)
assert client_config.clusters is not clusters
assert client_config.clusters['cluster1'] is cluster_cluster_config
def test_client_config_deserialize_public_key():
input_json = {
'clusters': {
'cluster1': {
'host': 'abc',
'user': '******',
'port': 22,
'auth': 'PUBLIC_KEY',
'key': '/home/user/.ssh/id_rsa',
'installKey': False,
'disableSshd': False,
'setupActions': {
'jupyter': ['echo a'],
'dask': []
},
'scratch': '$HOME',
'retries': {},
'useJupyterLab': True
}
},
'logLevel': INFO
}
client_config = ClientConfig(
clusters={
'cluster1': VALID_CLIENT_CLUSTER_CONFIG_WITH_PUBLIC_KEY_AUTH
})
assert deserialize_client_config_from_json(input_json) == client_config
def test_client_config_serialize_public_key():
client_config = ClientConfig(
clusters={
'cluster1': VALID_CLIENT_CLUSTER_CONFIG_WITH_PUBLIC_KEY_AUTH
})
expected_json = {
'clusters': {
'cluster1': {
'host': 'abc',
'user': '******',
'port': 22,
'auth': 'PUBLIC_KEY',
'key': '/home/user/.ssh/id_rsa',
'installKey': False,
'disableSshd': False,
'setupActions': {
'jupyter': ['echo a'],
'dask': []
},
'scratch': '$HOME',
'notebookDefaults': {},
'retries': DEFAULT_RETRIES_JSON,
'useJupyterLab': True
}
},
'logLevel': INFO
}
assert serialize_client_config_to_json(client_config) == expected_json
def test_client_config_deserialize():
LoggerProvider().log_level = DEBUG
input_json = {
'clusters': {
'cluster1': {
'host': 'abc',
'user': '******',
'port': 22,
'auth': 'ASK',
'key': None,
'installKey': True,
'disableSshd': False,
'setupActions': {
'jupyter': [],
'dask': []
},
'scratch': '$HOME',
'retries': {}
}
},
'logLevel': DEBUG
}
client_config = ClientConfig(clusters={
'cluster1':
ClusterConfigImpl(host='abc',
user='******',
port=22,
auth=AuthMethod.ASK)
},
log_level=DEBUG)
deserialized = deserialize_client_config_from_json(input_json)
pprint([i.__dict__ for i in deserialized.clusters.values()])
assert deserialize_client_config_from_json(input_json) == client_config
def test_client_config_serialize():
client_config = ClientConfig(
clusters={'cluster1': VALID_CLIENT_CLUSTER_CONFIG}, log_level=INFO)
expected_json = {
'clusters': {
'cluster1': {
'host': 'abc',
'user': '******',
'port': 22,
'auth': 'ASK',
'key': None,
'installKey': True,
'disableSshd': False,
'setupActions': {
'jupyter': [],
'dask': []
},
'scratch': '$HOME',
'notebookDefaults': {},
'retries': DEFAULT_RETRIES_JSON,
'useJupyterLab': True
}
},
'logLevel': INFO
}
assert serialize_client_config_to_json(client_config) == expected_json
def test_client_config_create_empty_and_add_cluster():
client_config = ClientConfig()
assert client_config.clusters == {}
cluster_cluster_config = VALID_CLIENT_CLUSTER_CONFIG
with pytest.raises(ValueError):
client_config.add_cluster(' Illegal Cluster Name',
cluster_cluster_config)
assert client_config.clusters == {}
client_config.add_cluster('cluster1', cluster_cluster_config)
assert client_config.clusters['cluster1'] is cluster_cluster_config
with pytest.raises(ValueError):
client_config.add_cluster(
'cluster1',
ClusterConfigImpl(host='another',
port=22,
user='******',
auth=AuthMethod.ASK))
assert client_config.clusters['cluster1'] is cluster_cluster_config
def test_create_environment():
"""Creating the environment and adding a cluster."""
cluster1 = ClusterConfigImpl(host='host1',
port=22,
user='******',
auth=AuthMethod.ASK)
config = ClientConfig(clusters={'cluster1': cluster1})
environment = EnvironmentImpl(config=config)
assert environment.config is config
assert len(environment.config.clusters) == 1
assert len(environment.clusters) == 1
assert isinstance(environment.clusters['cluster1'], Cluster)
cluster2 = ClusterConfigImpl(host='host2',
port=22,
user='******',
auth=AuthMethod.ASK)
environment.add_cluster(name='cluster2', config=cluster2)
assert len(environment.config.clusters) == 2
assert len(environment.clusters) == 2
assert isinstance(environment.clusters['cluster2'], Cluster)
def check_able_to_pull_environment(user: str,
remote_environment_upload_path: str,
remote_environment_pull_path: Optional[
str]):
with ExitStack() as stack:
stack.enter_context(disable_pytest_stdin())
stack.enter_context(reset_environment(user))
stack.enter_context(set_password(get_test_user_password(user)))
stack.enter_context(remove_remote_file(
user=user,
path=remote_environment_upload_path))
cluster = show_cluster(name=TEST_CLUSTER)
cluster.config.key = None
cluster.config.install_key = False
assert len(show_clusters()) == 1
fake_cluster = 'fake cluster'
remote_environment = EnvironmentImpl(
config=ClientConfig(
clusters={
TEST_CLUSTER: ClusterConfigImpl(
host=get_testing_host(),
port=get_testing_port(),
user=user,
auth=AuthMethod.ASK,
key='key_remote',
install_key=True),
fake_cluster: ClusterConfigImpl(
host='fakehost',
port=2,
user=user,
auth=AuthMethod.ASK,
key='key_remote',
install_key=True)}))
remote_environment_serialized = serialize_environment(
environment=remote_environment)
node = cluster.get_access_node()
assert isinstance(node, NodeInternal)
put_file_on_node(node=node,
remote_path=remote_environment_upload_path,
contents=remote_environment_serialized)
pull_environment(cluster=cluster,
path=remote_environment_pull_path)
assert len(show_clusters()) == 2
# Local key was kept unchanged.
assert show_cluster(TEST_CLUSTER).config == ClusterConfigImpl(
host=get_testing_host(),
port=get_testing_port(),
user=user,
auth=AuthMethod.ASK,
key=None,
install_key=False)
# New cluster was sanitized
assert show_cluster(fake_cluster).config == ClusterConfigImpl(
host='fakehost',
port=2,
user=user,
auth=AuthMethod.ASK,
key=None,
install_key=True)
def check_able_to_merge_push_environment(
user: str, remote_environment_upload_path: str,
remote_environment_push_path: Optional[str]):
with ExitStack() as stack:
stack.enter_context(disable_pytest_stdin())
stack.enter_context(reset_environment(user))
stack.enter_context(set_password(get_test_user_password(user)))
stack.enter_context(
remove_remote_file(user=user, path=remote_environment_upload_path))
cluster = show_cluster(name=TEST_CLUSTER)
cluster.config.key = None
cluster.config.install_key = False
cluster.config.retries[Retry.PORT_INFO] = \
get_default_retries_heavy_load()[Retry.PORT_INFO]
assert len(show_clusters()) == 1
node = cluster.get_access_node()
assert isinstance(node, NodeInternal)
# Upload an environment to be merged on push.
initial_remote_environment = EnvironmentImpl(config=ClientConfig(
clusters={
TEST_CLUSTER:
ClusterConfigImpl(host=get_testing_host(),
port=123,
user=user,
auth=AuthMethod.ASK,
key='key_remote',
install_key=True)
}))
initial_remote_environment_serialized = serialize_environment(
environment=initial_remote_environment)
put_file_on_node(node=node,
remote_path=remote_environment_upload_path,
contents=initial_remote_environment_serialized)
# Modify current environment.
fake_cluster = 'fake cluster'
add_cluster(name=fake_cluster,
host='fakehost',
port=2,
user=user,
auth=AuthMethod.ASK,
key='key_local',
install_key=False)
# Push with merge.
push_environment(cluster=cluster, path=remote_environment_push_path)
remote_environment_after_push_serialized = \
get_file_from_node(node=node,
remote_path=remote_environment_upload_path)
remote_environment_after_push = deserialize_environment(
text=remote_environment_after_push_serialized)
# Remote key was kept unchanged.
remote_clusters = remote_environment_after_push.config.clusters
assert len(remote_clusters) == 2
assert remote_clusters[TEST_CLUSTER].__dict__ == ClusterConfigImpl(
host=get_testing_host(),
port=get_testing_port(),
user=user,
auth=AuthMethod.ASK,
key='key_remote',
install_key=True,
retries=get_default_retries_heavy_load()).__dict__
# New cluster was sanitized
assert remote_clusters[fake_cluster].__dict__ == ClusterConfigImpl(
host='fakehost',
port=2,
user=user,
auth=AuthMethod.ASK,
key=None,
install_key=True).__dict__
def test_merge_environments():
auth = AuthMethod.PUBLIC_KEY
environment_1 = EnvironmentImpl(
config=ClientConfig(clusters={'a': ClusterConfigImpl(host='localhost1',
port=1,
user='******',
auth=auth,
key='key1'),
'b': ClusterConfigImpl(host='localhost2',
port=2,
user='******',
auth=auth,
key='key1'),
'c': ClusterConfigImpl(host='localhost3',
port=3,
user='******',
auth=auth)},
log_level=10))
environment_2 = EnvironmentImpl(
config=ClientConfig(clusters={'b': ClusterConfigImpl(host='localhost4',
port=4,
user='******',
auth=auth,
key='key2'),
'c': ClusterConfigImpl(host='localhost5',
port=5,
user='******',
auth=auth),
'd': ClusterConfigImpl(host='localhost6',
port=6,
user='******',
auth=auth,
key='key2')},
log_level=20))
merged_1_2 = merge_environments(local=environment_1,
remote=environment_2)
# Kept some fields from 1 after merge, replaced rest with 2.
# Sanitized some fields for new clusters from 2.
assert merged_1_2 == EnvironmentImpl(
config=ClientConfig(clusters={'a': ClusterConfigImpl(host='localhost1',
port=1,
user='******',
auth=auth,
key='key1'),
'b': ClusterConfigImpl(host='localhost4',
port=4,
user='******',
auth=auth,
key='key1'),
'c': ClusterConfigImpl(host='localhost5',
port=5,
user='******',
auth=auth),
'd': ClusterConfigImpl(host='localhost6',
port=6,
user='******',
auth=auth,
key=None)},
log_level=10))
merged_2_1 = merge_environments(local=environment_2,
remote=environment_1)
# Kept some fields from 2 after merge, replaced rest with 1.
# Sanitized some fields for new clusters from 1.
assert merged_2_1 == EnvironmentImpl(
config=ClientConfig(clusters={'a': ClusterConfigImpl(host='localhost1',
port=1,
user='******',
auth=auth,
key=None),
'b': ClusterConfigImpl(host='localhost2',
port=2,
user='******',
auth=auth,
key='key2'),
'c': ClusterConfigImpl(host='localhost3',
port=3,
user='******',
auth=auth),
'd': ClusterConfigImpl(host='localhost6',
port=6,
user='******',
auth=auth,
key='key2')},
log_level=20))
def test_client_config_validation_is_used():
cluster_cluster_config = VALID_CLIENT_CLUSTER_CONFIG
with pytest.raises(ValueError):
ClientConfig({' Illegal Cluster Name': cluster_cluster_config})