def test_import_local_settings_without_syspath(self, log_mock):
"""
Tests that an ImportError is raised in import_local_settings
if there is no airflow_local_settings module on the syspath.
"""
from airflow import settings
settings.import_local_settings()
log_mock.assert_called_with("Failed to import airflow_local_settings.", exc_info=True)
def test_custom_policy(self):
with SettingsContext(SETTINGS_FILE_CUSTOM_POLICY,
"airflow_local_settings"):
from airflow import settings
settings.import_local_settings()
task_instance = MagicMock()
task_instance.owner = 'airflow'
with self.assertRaises(AirflowClusterPolicyViolation):
settings.task_must_have_owners(task_instance) # pylint: disable=no-member
def test_import_with_dunder_all_not_specified(self):
"""
Tests that if __all__ is specified in airflow_local_settings,
only module attributes specified within are imported.
"""
with SettingsContext(SETTINGS_FILE_POLICY_WITH_DUNDER_ALL, "airflow_local_settings"):
from airflow import settings
settings.import_local_settings() # pylint: ignore
with self.assertRaises(AttributeError):
settings.not_policy()
def test_policy_function(self):
"""
Tests that task instances are mutated by the policy
function in airflow_local_settings.
"""
with SettingsContext(SETTINGS_FILE_POLICY, "airflow_local_settings"):
from airflow import settings
settings.import_local_settings()
task_instance = MagicMock()
settings.test_policy(task_instance) # pylint: disable=no-member
assert task_instance.run_as_user == "myself"
def test_pod_mutation_hook(self):
"""
Tests that pods are mutated by the pod_mutation_hook
function in airflow_local_settings.
"""
with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK, "airflow_local_settings"):
from airflow import settings
settings.import_local_settings() # pylint: ignore
pod = Pod(image="ubuntu", envs={}, cmds=['echo "1"'])
settings.pod_mutation_hook(pod)
assert pod.namespace == 'airflow-tests'
def test_import_with_dunder_all(self):
"""
Tests that if __all__ is specified in airflow_local_settings,
only module attributes specified within are imported.
"""
with SettingsContext(SETTINGS_FILE_POLICY_WITH_DUNDER_ALL, "airflow_local_settings"):
from airflow import settings
settings.import_local_settings() # pylint: ignore
task_instance = MagicMock()
settings.test_policy(task_instance)
assert task_instance.run_as_user == "myself"
def test_pod_mutation_hook(self):
"""
Tests that pods are mutated by the pod_mutation_hook
function in airflow_local_settings.
"""
with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK,
"airflow_local_settings"):
from airflow import settings
settings.import_local_settings()
pod = MagicMock()
settings.pod_mutation_hook(pod)
assert pod.namespace == 'airflow-tests'
def test_pod_mutation_v1_pod(self):
with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK_V1_POD,
"airflow_local_settings"):
from airflow import settings
settings.import_local_settings() # pylint: ignore
from airflow.kubernetes.pod_launcher import PodLauncher
self.mock_kube_client = Mock()
self.pod_launcher = PodLauncher(kube_client=self.mock_kube_client)
pod = pod_generator.PodGenerator(image="myimage",
cmds=["foo"],
namespace="baz",
volume_mounts=[{
"name": "foo",
"mountPath": "/mnt",
"subPath": "/",
"readOnly": True
}],
volumes=[{
"name": "foo"
}]).gen_pod()
sanitized_pod_pre_mutation = api_client.sanitize_for_serialization(
pod)
self.assertEqual(
sanitized_pod_pre_mutation, {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'namespace': 'baz'
},
'spec': {
'containers': [{
'args': [],
'command': ['foo'],
'env': [],
'envFrom': [],
'image':
'myimage',
'name':
'base',
'ports': [],
'volumeMounts': [{
'mountPath': '/mnt',
'name': 'foo',
'readOnly': True,
'subPath': '/'
}]
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'volumes': [{
'name': 'foo'
}]
}
})
# Apply Pod Mutation Hook
pod = self.pod_launcher._mutate_pod_backcompat(pod)
sanitized_pod_post_mutation = api_client.sanitize_for_serialization(
pod)
self.assertEqual(
sanitized_pod_post_mutation, {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'namespace': 'airflow-tests'
},
'spec': {
'containers': [{
'args': [],
'command': ['foo'],
'env': [{
'name': 'TEST_USER',
'value': 'ADMIN'
}],
'envFrom': [],
'image':
'test-image',
'name':
'base',
'ports': [{
'containerPort': 8080
}, {
'containerPort': 8081
}],
'volumeMounts': [{
'mountPath': '/mnt',
'name': 'foo',
'readOnly': True,
'subPath': '/'
}, {
'mountPath': '/opt/airflow/secrets/',
'name': 'airflow-secrets-mount',
'readOnly': True
}]
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'volumes': [{
'name': 'foo'
}, {
'name': 'airflow-secrets-mount',
'secret': {
'secretName': 'airflow-test-secrets'
}
}]
}
})
def test_pod_mutation_to_k8s_pod(self):
with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK,
"airflow_local_settings"):
from airflow import settings
settings.import_local_settings() # pylint: ignore
from airflow.kubernetes.pod_launcher import PodLauncher
self.mock_kube_client = Mock()
self.pod_launcher = PodLauncher(kube_client=self.mock_kube_client)
init_container = k8s.V1Container(name="init-container",
volume_mounts=[
k8s.V1VolumeMount(
mount_path="/tmp",
name="init-secret")
])
pod = pod_generator.PodGenerator(
image="foo",
name="bar",
namespace="baz",
image_pull_policy="Never",
init_containers=[init_container],
cmds=["foo"],
args=["/bin/sh", "-c", "touch /tmp/healthy"],
tolerations=[{
'effect': 'NoSchedule',
'key': 'static-pods',
'operator': 'Equal',
'value': 'true'
}],
volume_mounts=[{
"name": "foo",
"mountPath": "/mnt",
"subPath": "/",
"readOnly": True
}],
security_context=k8s.V1PodSecurityContext(fs_group=0,
run_as_user=1),
volumes=[k8s.V1Volume(name="foo")]).gen_pod()
sanitized_pod_pre_mutation = api_client.sanitize_for_serialization(
pod)
self.assertEqual(
sanitized_pod_pre_mutation,
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': mock.ANY,
'namespace': 'baz'
},
'spec': {
'containers': [{
'args': ['/bin/sh', '-c', 'touch /tmp/healthy'],
'command': ['foo'],
'env': [],
'envFrom': [],
'image':
'foo',
'imagePullPolicy':
'Never',
'name':
'base',
'ports': [],
'volumeMounts': [{
'mountPath': '/mnt',
'name': 'foo',
'readOnly': True,
'subPath': '/'
}]
}],
'initContainers': [{
'name':
'init-container',
'volumeMounts': [{
'mountPath': '/tmp',
'name': 'init-secret'
}]
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'tolerations': [{
'effect': 'NoSchedule',
'key': 'static-pods',
'operator': 'Equal',
'value': 'true'
}],
'volumes': [{
'name': 'foo'
}],
'securityContext': {
'fsGroup': 0,
'runAsUser': 1
}
}
},
)
# Apply Pod Mutation Hook
pod = self.pod_launcher._mutate_pod_backcompat(pod)
sanitized_pod_post_mutation = api_client.sanitize_for_serialization(
pod)
self.assertEqual(
sanitized_pod_post_mutation, {
"apiVersion": "v1",
"kind": "Pod",
'metadata': {
'labels': {
'test_label': 'test_value'
},
'name': mock.ANY,
'namespace': 'airflow-tests'
},
'spec': {
'affinity': {
'nodeAffinity': {
'requiredDuringSchedulingIgnoredDuringExecution':
{
'nodeSelectorTerms': [{
'matchExpressions': [{
'key': 'test/dynamic-pods',
'operator': 'In',
'values': ['true']
}]
}]
}
}
},
'containers': [{
'args': ['/bin/sh', '-c', 'touch /tmp/healthy2'],
'command': ['foo'],
'env': [{
'name': 'TEST_USER',
'value': 'ADMIN'
}],
'image':
'my_image',
'imagePullPolicy':
'Never',
'name':
'base',
'ports': [{
'containerPort': 8080
}, {
'containerPort': 8081
}],
'resources': {
'limits': {
'nvidia.com/gpu': '200G'
},
'requests': {
'cpu': '200Mi',
'memory': '2G'
}
},
'volumeMounts': [{
'mountPath': '/mnt',
'name': 'foo',
'readOnly': True,
'subPath': '/'
}, {
'mountPath': '/opt/airflow/secrets/',
'name': 'airflow-secrets-mount',
'readOnly': True
}]
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'initContainers': [{
'name':
'init-container',
'securityContext': {
'runAsGroup': 50000,
'runAsUser': 50000
},
'volumeMounts': [{
'mountPath': '/tmp',
'name': 'init-secret'
}]
}],
'tolerations': [{
'effect': 'NoSchedule',
'key': 'static-pods',
'operator': 'Equal',
'value': 'true'
}, {
'effect': 'NoSchedule',
'key': 'dynamic-pods',
'operator': 'Equal',
'value': 'true'
}],
'volumes': [
{
'name': 'airflow-secrets-mount',
'secret': {
'secretName': 'airflow-test-secrets'
}
},
{
'name': 'bar'
},
{
'name': 'foo'
},
],
'securityContext': {
'runAsUser': 1
}
}
})