[airflow] branch master updated: Add tests for Custom cluster policy (#11381)

[po...@apache.org]

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

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


The following commit(s) were added to refs/heads/master by this push:
     new 6fe020e  Add tests for Custom cluster policy (#11381)
6fe020e is described below

commit 6fe020e105531dd5a7097d8875eac0f317045298
Author: Jarek Potiuk <ja...@polidea.com>
AuthorDate: Sat Oct 10 00:57:10 2020 +0200

    Add tests for Custom cluster policy (#11381)
    
    The custom ClusterPolicyViolation has been added in #10282
    This one adds more comprehensive test to it.
    
    Co-authored-by: Jacob Ferriero <jf...@google.com>
---
 tests/test_local_settings.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tests/test_local_settings.py b/tests/test_local_settings.py
index f05e8f8..09e03f8 100644
--- a/tests/test_local_settings.py
+++ b/tests/test_local_settings.py
@@ -22,6 +22,8 @@ import tempfile
 import unittest
 from unittest.mock import MagicMock, call
 
+from airflow.exceptions import AirflowClusterPolicyViolation
+
 SETTINGS_FILE_POLICY = """
 def test_policy(task_instance):
     task_instance.run_as_user = "myself"
@@ -42,6 +44,18 @@ def pod_mutation_hook(pod):
     pod.namespace = 'airflow-tests'
 """
 
+SETTINGS_FILE_CUSTOM_POLICY = """
+from airflow.models.baseoperator import BaseOperator
+from airflow.exceptions import AirflowClusterPolicyViolation
+
+def task_must_have_owners(task: BaseOperator):
+    if not task.owner or task.owner.lower() == "airflow":
+        raise AirflowClusterPolicyViolation(
+            f'''Task must have non-None non-'airflow' owner.
+            Current value: {task.owner}'''
+        )
+"""
+
 
 class SettingsContext:
     def __init__(self, content: str, module_name: str):
@@ -149,3 +163,13 @@ class TestLocalSettings(unittest.TestCase):
             settings.pod_mutation_hook(pod)
 
             assert pod.namespace == 'airflow-tests'
+
+    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