def test_spark_task_service_env(self): key = secret.make_secret_key() service_env = dict([(key, '{} value'.format(key)) for key in SparkTaskService.SERVICE_ENV_KEYS]) service_env.update({"other": "value"}) with os_environ(service_env): service = SparkTaskService(1, key, None) client = SparkTaskClient(1, service.addresses(), key, 3) with tempdir() as d: file = '{}/env'.format(d) command = "env | grep -v '^PWD='> {}".format(file) command_env = {"test": "value"} try: client.run_command(command, command_env) client.wait_for_command_termination() finally: service.shutdown() with open(file) as f: env = sorted([line.strip() for line in f.readlines()]) expected = [ 'HADOOP_TOKEN_FILE_LOCATION=HADOOP_TOKEN_FILE_LOCATION value', 'test=value' ] self.assertEqual(env, expected)
def spark_task_service(index, key=None, nics=None, match_intf=False, minimum_command_lifetime_s=0, verbose=2): key = key or secret.make_secret_key() task = SparkTaskService(index, key, nics, minimum_command_lifetime_s, verbose) client = SparkTaskClient(index, task.addresses(), key, verbose, match_intf) try: yield task, client, key finally: task.shutdown()