Python worker_configuration.WorkerConfiguration类代码示例

    def test_worker_git_dags(self):
        # Tests persistence volume config created when `git_repo` is set
        self.kube_config.dags_volume_claim = None
        self.kube_config.dags_volume_host = None
        self.kube_config.dags_folder = '/usr/local/airflow/dags'
        self.kube_config.worker_dags_folder = '/usr/local/airflow/dags'

        self.kube_config.git_sync_container_repository = 'gcr.io/google-containers/git-sync-amd64'
        self.kube_config.git_sync_container_tag = 'v2.0.5'
        self.kube_config.git_sync_container = 'gcr.io/google-containers/git-sync-amd64:v2.0.5'
        self.kube_config.git_sync_init_container_name = 'git-sync-clone'
        self.kube_config.git_subpath = 'dags_folder'
        self.kube_config.git_sync_root = '/git'
        self.kube_config.git_dags_folder_mount_point = '/usr/local/airflow/dags/repo/dags_folder'

        worker_config = WorkerConfiguration(self.kube_config)
        volumes, volume_mounts = worker_config.init_volumes_and_mounts()

        dag_volume = [volume for volume in volumes.values() if volume['name'] == 'airflow-dags']
        dag_volume_mount = [mount for mount in volume_mounts.values() if mount['name'] == 'airflow-dags']

        self.assertTrue('emptyDir' in dag_volume[0])
        self.assertEqual(self.kube_config.git_dags_folder_mount_point, dag_volume_mount[0]['mountPath'])
        self.assertTrue(dag_volume_mount[0]['readOnly'])

        init_container = worker_config._get_init_containers(volume_mounts)[0]
        init_container_volume_mount = [mount for mount in init_container['volumeMounts']
                                       if mount['name'] == 'airflow-dags']

        self.assertEqual('git-sync-clone', init_container['name'])
        self.assertEqual('gcr.io/google-containers/git-sync-amd64:v2.0.5', init_container['image'])
        self.assertEqual(1, len(init_container_volume_mount))
        self.assertFalse(init_container_volume_mount[0]['readOnly'])

    def test_make_pod_git_sync_ssh_without_known_hosts(self):
        # Tests the pod created with git-sync SSH authentication option is correct without known hosts
        self.kube_config.airflow_configmap = 'airflow-configmap'
        self.kube_config.git_ssh_key_secret_name = 'airflow-secrets'
        self.kube_config.dags_volume_claim = None
        self.kube_config.dags_volume_host = None
        self.kube_config.dags_in_image = None
        self.kube_config.worker_fs_group = None

        worker_config = WorkerConfiguration(self.kube_config)
        kube_executor_config = KubernetesExecutorConfig(annotations=[],
                                                        volumes=[],
                                                        volume_mounts=[])

        pod = worker_config.make_pod("default", str(uuid.uuid4()), "test_pod_id", "test_dag_id",
                                     "test_task_id", str(datetime.utcnow()), 1, "bash -c 'ls /'",
                                     kube_executor_config)

        init_containers = worker_config._get_init_containers()
        git_ssh_key_file = next((x['value'] for x in init_containers[0]['env']
                                if x['name'] == 'GIT_SSH_KEY_FILE'), None)
        volume_mount_ssh_key = next((x['mountPath'] for x in init_containers[0]['volumeMounts']
                                    if x['name'] == worker_config.git_sync_ssh_secret_volume_name),
                                    None)
        self.assertTrue(git_ssh_key_file)
        self.assertTrue(volume_mount_ssh_key)
        self.assertEqual(65533, pod.security_context['fsGroup'])
        self.assertEqual(git_ssh_key_file,
                         volume_mount_ssh_key,
                         'The location where the git ssh secret is mounted'
                         ' needs to be the same as the GIT_SSH_KEY_FILE path')

    def test_init_environment_using_git_sync_user_with_known_hosts(self):
        # Tests the init environment created with git-sync User authentication option is correct
        # with known hosts file
        self.kube_config.airflow_configmap = 'airflow-configmap'
        self.kube_config.git_user = 'git_user'
        self.kube_config.git_password = 'git_password'
        self.kube_config.git_ssh_known_hosts_configmap_name = 'airflow-configmap'
        self.kube_config.git_ssh_key_secret_name = None
        self.kube_config.dags_volume_claim = None
        self.kube_config.dags_volume_host = None
        self.kube_config.dags_in_image = None

        worker_config = WorkerConfiguration(self.kube_config)
        init_containers = worker_config._get_init_containers()

        self.assertTrue(init_containers)  # check not empty
        env = init_containers[0]['env']

        self.assertFalse({'name': 'GIT_SSH_KEY_FILE', 'value': '/etc/git-secret/ssh'} in env)
        self.assertTrue({'name': 'GIT_SYNC_USERNAME', 'value': 'git_user'} in env)
        self.assertTrue({'name': 'GIT_SYNC_PASSWORD', 'value': 'git_password'} in env)
        self.assertTrue({'name': 'GIT_KNOWN_HOSTS', 'value': 'true'} in env)
        self.assertTrue({'name': 'GIT_SSH_KNOWN_HOSTS_FILE',
                        'value': '/etc/git-secret/known_hosts'} in env)
        self.assertFalse({'name': 'GIT_SYNC_SSH', 'value': 'true'} in env)

    def test_get_secrets(self):
        # Test when secretRef is None and kube_secrets is not empty
        self.kube_config.kube_secrets = {
            'AWS_SECRET_KEY': 'airflow-secret=aws_secret_key',
            'POSTGRES_PASSWORD': 'airflow-secret=postgres_credentials'
        }
        self.kube_config.env_from_secret_ref = None
        worker_config = WorkerConfiguration(self.kube_config)
        secrets = worker_config._get_secrets()
        secrets.sort(key=lambda secret: secret.deploy_target)
        expected = [
            Secret('env', 'AWS_SECRET_KEY', 'airflow-secret', 'aws_secret_key'),
            Secret('env', 'POSTGRES_PASSWORD', 'airflow-secret', 'postgres_credentials')
        ]
        self.assertListEqual(expected, secrets)

        # Test when secret is not empty and kube_secrets is empty dict
        self.kube_config.kube_secrets = {}
        self.kube_config.env_from_secret_ref = 'secret_a,secret_b'
        worker_config = WorkerConfiguration(self.kube_config)
        secrets = worker_config._get_secrets()
        expected = [
            Secret('env', None, 'secret_a'),
            Secret('env', None, 'secret_b')
        ]
        self.assertListEqual(expected, secrets)

    def test_worker_environment_no_dags_folder(self):
        self.kube_config.airflow_configmap = ''
        self.kube_config.git_dags_folder_mount_point = ''
        self.kube_config.dags_folder = ''
        worker_config = WorkerConfiguration(self.kube_config)
        env = worker_config._get_environment()

        self.assertNotIn('AIRFLOW__CORE__DAGS_FOLDER', env)

    def test_worker_environment_when_dags_folder_specified(self):
        dags_folder = '/workers/path/to/dags'
        self.kube_config.worker_dags_folder = dags_folder

        worker_config = WorkerConfiguration(self.kube_config)
        env = worker_config._get_environment()

        self.assertEqual(dags_folder, env['AIRFLOW__CORE__DAGS_FOLDER'])

 def test_worker_configuration_no_subpaths(self):
     worker_config = WorkerConfiguration(self.kube_config)
     volumes, volume_mounts = worker_config.init_volumes_and_mounts()
     for volume_or_mount in volumes + volume_mounts:
         if volume_or_mount['name'] != 'airflow-config':
             self.assertNotIn(
                 'subPath', volume_or_mount,
                 "subPath shouldn't be defined"
             )

    def test_worker_environment_when_dags_folder_specified(self):
        self.kube_config.airflow_configmap = 'airflow-configmap'
        self.kube_config.git_dags_folder_mount_point = ''
        dags_folder = '/workers/path/to/dags'
        self.kube_config.dags_folder = dags_folder

        worker_config = WorkerConfiguration(self.kube_config)
        env = worker_config._get_environment()

        self.assertEqual(dags_folder, env['AIRFLOW__CORE__DAGS_FOLDER'])

 def test_worker_configuration_no_subpaths(self):
     worker_config = WorkerConfiguration(self.kube_config)
     volumes, volume_mounts = worker_config._get_volumes_and_mounts()
     volumes_list = [value for value in volumes.values()]
     volume_mounts_list = [value for value in volume_mounts.values()]
     for volume_or_mount in volumes_list + volume_mounts_list:
         if volume_or_mount['name'] != 'airflow-config':
             self.assertNotIn(
                 'subPath', volume_or_mount,
                 "subPath shouldn't be defined"
             )

    def test_worker_pvc_dags(self):
        # Tests persistence volume config created when `dags_volume_claim` is set
        self.kube_config.dags_volume_claim = 'airflow-dags'

        worker_config = WorkerConfiguration(self.kube_config)
        volumes, volume_mounts = worker_config.init_volumes_and_mounts()

        dag_volume = [volume for volume in volumes if volume['name'] == 'airflow-dags']
        dag_volume_mount = [mount for mount in volume_mounts if mount['name'] == 'airflow-dags']

        self.assertEqual('airflow-dags', dag_volume[0]['persistentVolumeClaim']['claimName'])
        self.assertEqual(1, len(dag_volume_mount))

    def test_get_configmaps(self):
        # Test when configmap is empty
        self.kube_config.env_from_configmap_ref = ''
        worker_config = WorkerConfiguration(self.kube_config)
        configmaps = worker_config._get_configmaps()
        self.assertListEqual([], configmaps)

        # test when configmap is not empty
        self.kube_config.env_from_configmap_ref = 'configmap_a,configmap_b'
        worker_config = WorkerConfiguration(self.kube_config)
        configmaps = worker_config._get_configmaps()
        self.assertListEqual(['configmap_a', 'configmap_b'], configmaps)

    def test_worker_environment_dags_folder_using_git_sync(self):
        self.kube_config.airflow_configmap = 'airflow-configmap'
        self.kube_config.git_sync_dest = 'repo'
        self.kube_config.git_subpath = 'dags'
        self.kube_config.git_dags_folder_mount_point = '/workers/path/to/dags'

        dags_folder = '{}/{}/{}'.format(self.kube_config.git_dags_folder_mount_point,
                                        self.kube_config.git_sync_dest,
                                        self.kube_config.git_subpath)

        worker_config = WorkerConfiguration(self.kube_config)
        env = worker_config._get_environment()

        self.assertEqual(dags_folder, env['AIRFLOW__CORE__DAGS_FOLDER'])

    def test_worker_container_dags(self):
        # Tests that the 'airflow-dags' persistence volume is NOT created when `dags_in_image` is set
        self.kube_config.dags_in_image = True

        worker_config = WorkerConfiguration(self.kube_config)
        volumes, volume_mounts = worker_config.init_volumes_and_mounts()

        dag_volume = [volume for volume in volumes.values() if volume['name'] == 'airflow-dags']
        dag_volume_mount = [mount for mount in volume_mounts.values() if mount['name'] == 'airflow-dags']

        init_containers = worker_config._get_init_containers(volume_mounts)

        self.assertEqual(0, len(dag_volume))
        self.assertEqual(0, len(dag_volume_mount))
        self.assertEqual(0, len(init_containers))

    def test_worker_with_subpaths(self):
        self.kube_config.dags_volume_subpath = 'dags'
        self.kube_config.logs_volume_subpath = 'logs'
        worker_config = WorkerConfiguration(self.kube_config)
        volumes, volume_mounts = worker_config.init_volumes_and_mounts()

        for volume in volumes:
            self.assertNotIn(
                'subPath', volume,
                "subPath isn't valid configuration for a volume"
            )

        for volume_mount in volume_mounts:
            if volume_mount['name'] != 'airflow-config':
                self.assertIn(
                    'subPath', volume_mount,
                    "subPath should've been passed to volumeMount configuration"
                )

    def test_worker_generate_dag_volume_mount_path(self):
        self.kube_config.git_dags_folder_mount_point = '/root/airflow/git/dags'
        self.kube_config.dags_folder = '/root/airflow/dags'
        worker_config = WorkerConfiguration(self.kube_config)

        self.kube_config.dags_volume_claim = 'airflow-dags'
        self.kube_config.dags_volume_host = ''
        dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
        self.assertEqual(dag_volume_mount_path, self.kube_config.dags_folder)

        self.kube_config.dags_volume_claim = ''
        self.kube_config.dags_volume_host = '/host/airflow/dags'
        dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
        self.assertEqual(dag_volume_mount_path, self.kube_config.dags_folder)

        self.kube_config.dags_volume_claim = ''
        self.kube_config.dags_volume_host = ''
        dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
        self.assertEqual(dag_volume_mount_path,
                         self.kube_config.git_dags_folder_mount_point)