[airflow] branch main updated: Adding typing for KPO SCC objects (#33381)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 46ffc25995 Adding typing for KPO SCC objects (#33381)
46ffc25995 is described below

commit 46ffc259956d3782bb0c09e782fca634ff4215b9
Author: Amogh Desai <am...@gmail.com>
AuthorDate: Mon Aug 14 17:20:52 2023 +0530

    Adding typing for KPO SCC objects (#33381)
    
    * Adding typing for KPO SCC objects
    
    
    ---------
    
    Co-authored-by: Hussein Awala <hu...@awala.fr>
    Co-authored-by: Wei Lee <we...@gmail.com>
    Co-authored-by: John Brandborg <jo...@protonmail.com>
---
 airflow/providers/cncf/kubernetes/operators/pod.py    |  4 ++--
 kubernetes_tests/test_kubernetes_pod_operator.py      | 12 +++++++-----
 tests/providers/cncf/kubernetes/operators/test_pod.py |  6 ++----
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
index d72b2fec0c..a3958b6237 100644
--- a/airflow/providers/cncf/kubernetes/operators/pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -296,8 +296,8 @@ class KubernetesPodOperator(BaseOperator):
         service_account_name: str | None = None,
         hostnetwork: bool = False,
         tolerations: list[k8s.V1Toleration] | None = None,
-        security_context: dict | None = None,
-        container_security_context: dict | None = None,
+        security_context: k8s.V1PodSecurityContext | dict | None = None,
+        container_security_context: k8s.V1SecurityContext | dict | None = None,
         dnspolicy: str | None = None,
         dns_config: k8s.V1PodDNSConfig | None = None,
         hostname: str | None = None,
diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py
index 5b94e2186c..3d7a2fd4ef 100644
--- a/kubernetes_tests/test_kubernetes_pod_operator.py
+++ b/kubernetes_tests/test_kubernetes_pod_operator.py
@@ -29,7 +29,7 @@ from uuid import uuid4
 
 import pendulum
 import pytest
-from kubernetes.client import V1EnvVar, models as k8s
+from kubernetes.client import V1EnvVar, V1PodSecurityContext, V1SecurityContext, models as k8s
 from kubernetes.client.api_client import ApiClient
 from kubernetes.client.rest import ApiException
 from pytest import param
@@ -514,7 +514,7 @@ class TestKubernetesPodOperatorSystem:
 
     @pytest.mark.parametrize("uid", [0, 1000])
     def test_run_as_user(self, uid, mock_get_connection):
-        security_context = {"runAsUser": uid}
+        security_context = V1PodSecurityContext(run_as_user=uid)
         name = str(uuid4())
         k = KubernetesPodOperator(
             namespace="default",
@@ -539,7 +539,7 @@ class TestKubernetesPodOperatorSystem:
 
     @pytest.mark.parametrize("gid", [0, 1000])
     def test_fs_group(self, gid, mock_get_connection):
-        security_context = {"fsGroup": gid}
+        security_context = V1PodSecurityContext(fs_group=gid)
         name = str(uuid4())
         k = KubernetesPodOperator(
             namespace="default",
@@ -563,7 +563,7 @@ class TestKubernetesPodOperatorSystem:
         assert pod.to_dict()["spec"]["security_context"]["fs_group"] == gid
 
     def test_disable_privilege_escalation(self, mock_get_connection):
-        container_security_context = {"allowPrivilegeEscalation": False}
+        container_security_context = V1SecurityContext(allow_privilege_escalation=False)
 
         k = KubernetesPodOperator(
             namespace="default",
@@ -579,7 +579,9 @@ class TestKubernetesPodOperatorSystem:
         context = create_context(k)
         k.execute(context)
         actual_pod = self.api_client.sanitize_for_serialization(k.pod)
-        self.expected_pod["spec"]["containers"][0]["securityContext"] = container_security_context
+        self.expected_pod["spec"]["containers"][0]["securityContext"] = {
+            "allowPrivilegeEscalation": container_security_context.allow_privilege_escalation
+        }
         assert self.expected_pod == actual_pod
 
     def test_faulty_image(self, mock_get_connection):
diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py
index b398ec2053..7b40103dac 100644
--- a/tests/providers/cncf/kubernetes/operators/test_pod.py
+++ b/tests/providers/cncf/kubernetes/operators/test_pod.py
@@ -23,7 +23,7 @@ from unittest.mock import MagicMock, patch
 
 import pendulum
 import pytest
-from kubernetes.client import ApiClient, V1PodStatus, models as k8s
+from kubernetes.client import ApiClient, V1PodSecurityContext, V1PodStatus, models as k8s
 from pytest import param
 from urllib3 import HTTPResponse
 from urllib3.packages.six import BytesIO
@@ -219,9 +219,7 @@ class TestKubernetesPodOperator:
         assert k.env_vars[0].name == "bartemplated"
 
     def test_security_context(self):
-        security_context = {
-            "runAsUser": 1245,
-        }
+        security_context = V1PodSecurityContext(run_as_user=1245)
         k = KubernetesPodOperator(
             security_context=security_context,
             task_id="task",