diff --git a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py index 213614c80c82e1..bf0cf81c0ceb5f 100644 --- a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py +++ b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py @@ -257,7 +257,7 @@ def __init__( # pylint: disable=too-many-arguments,too-many-locals else: self.node_selector = {} self.annotations = annotations or {} - self.affinity = convert_affinity(affinity) if affinity else k8s.V1Affinity() + self.affinity = convert_affinity(affinity) if affinity else {} self.k8s_resources = convert_resources(resources) if resources else {} self.config_file = config_file self.image_pull_secrets = convert_image_pull_secrets(image_pull_secrets) if image_pull_secrets else [] diff --git a/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py b/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py index a087fc9b0edf31..2da7e8b9dbb8f9 100644 --- a/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py @@ -334,6 +334,25 @@ def test_pod_template_file(self): labels: foo: bar spec: + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: kubernetes.io/role + operator: In + values: + - foo + - bar + preferredDuringSchedulingIgnoredDuringExecution: + - weight: 1 + preference: + matchExpressions: + - key: kubernetes.io/role + operator: In + values: + - foo + - bar containers: - name: base image: ubuntu:16.04 @@ -366,6 +385,35 @@ def test_pod_template_file(self): assert pod.metadata.namespace == "mynamespace" assert pod.spec.containers[0].image == "ubuntu:16.04" assert pod.spec.containers[0].command == ["something"] + affinity = { + 'node_affinity': { + 'preferred_during_scheduling_ignored_during_execution': [ + { + 'preference': { + 'match_expressions': [ + {'key': 'kubernetes.io/role', 'operator': 'In', 'values': ['foo', 'bar']} + ], + 'match_fields': None, + }, + 'weight': 1, + } + ], + 'required_during_scheduling_ignored_during_execution': { + 'node_selector_terms': [ + { + 'match_expressions': [ + {'key': 'kubernetes.io/role', 'operator': 'In', 'values': ['foo', 'bar']} + ], + 'match_fields': None, + } + ] + }, + }, + 'pod_affinity': None, + 'pod_anti_affinity': None, + } + + assert pod.spec.affinity.to_dict() == affinity # kwargs take precedence, however image = "some.custom.image:andtag"