def conf_vars(overrides):
original = {}
original_env_vars = {}
for (section, key), value in overrides.items():
env = conf._env_var_name(section, key)
if env in os.environ:
original_env_vars[env] = os.environ.pop(env)
if conf.has_option(section, key):
original[(section, key)] = conf.get(section, key)
else:
original[(section, key)] = None
if value is not None:
if not conf.has_section(section):
conf.add_section(section)
conf.set(section, key, value)
else:
conf.remove_option(section, key)
settings.configure_vars()
try:
yield
finally:
for (section, key), value in original.items():
if value is not None:
conf.set(section, key, value)
else:
conf.remove_option(section, key)
for env, value in original_env_vars.items():
os.environ[env] = value
settings.configure_vars()
def setUp(self):
if not conf.has_section("kerberos"):
conf.add_section("kerberos")
conf.set("kerberos", "keytab",
os.environ['KRB5_KTNAME'])
keytab_from_cfg = conf.get("kerberos", "keytab")
self.args = Namespace(keytab=keytab_from_cfg, principal=None, pid=None,
daemon=None, stdout=None, stderr=None, log_file=None)
def setUp(self):
try:
conf.add_section("api")
except Exception: # pylint:disable=broad-except
pass
conf.set("api", "auth_backend",
"airflow.api.auth.backend.kerberos_auth")
try:
conf.add_section("kerberos")
except Exception: # pylint:disable=broad-except
pass
conf.set("kerberos", "keytab", os.environ['KRB5_KTNAME'])
self.app, _ = application.create_app(testing=True)
def setUp(self):
try:
conf.add_section("atlas")
except AirflowConfigException:
pass
except DuplicateSectionError:
pass
conf.set("atlas", "username", "none")
conf.set("atlas", "password", "none")
conf.set("atlas", "host", "none")
conf.set("atlas", "port", "0")
self.atlas = AtlasBackend()
def __set_mocked_executor(self):
"""Mock ECS such that there's nothing wrong with anything"""
from airflow.configuration import conf
if not conf.has_section('batch'):
conf.add_section('batch')
conf.set('batch', 'region', 'us-west-1')
conf.set('batch', 'job_name', 'some-job-name')
conf.set('batch', 'job_queue', 'some-job-queue')
conf.set('batch', 'job_definition', 'some-job-def')
executor = AwsBatchExecutor()
executor.start()
# replace boto3 ecs client with mock
batch_mock = mock.Mock(spec=executor.batch)
submit_job_ret_val = {
'jobName': conf.get('batch', 'job_name'),
'jobId': 'ABC'
}
batch_mock.submit_job.return_value = submit_job_ret_val
executor.batch = batch_mock
self.executor = executor
def __set_mocked_executor(self):
"""Mock ECS such that there's nothing wrong with anything"""
from airflow.configuration import conf
if not conf.has_section('ecs_fargate'):
conf.add_section('ecs_fargate')
conf.set('ecs_fargate', 'region', 'us-west-1')
conf.set('ecs_fargate', 'cluster', 'some-ecs-cluster')
conf.set('ecs_fargate', 'task_definition', 'some-ecs-task-definition')
conf.set('ecs_fargate', 'container_name', 'some-ecs-container')
conf.set('ecs_fargate', 'launch_type', 'FARGATE')
executor = AwsEcsFargateExecutor()
executor.start()
# replace boto3 ecs client with mock
ecs_mock = mock.Mock(spec=executor.ecs)
run_task_ret_val = {
'tasks': [{'taskArn': '001'}],
'failures': []
}
ecs_mock.run_task.return_value = run_task_ret_val
executor.ecs = ecs_mock
self.executor = executor
