PrivacyScanData/Azure/test_azure_container_instance.py

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import json
import unittest
from unittest.mock import patch

from azure.mgmt.containerinstance.models import (
    Container,
    ContainerGroup,
    Logs,
    ResourceRequests,
    ResourceRequirements,
)

from airflow.models import Connection
from airflow.providers.microsoft.azure.hooks.container_instance import AzureContainerInstanceHook
from airflow.utils import db


class TestAzureContainerInstanceHook(unittest.TestCase):
    def setUp(self):
        db.merge_conn(
            Connection(
                conn_id='azure_container_instance_test',
                conn_type='azure_container_instances',
                login='login',
                password='key',
                extra=json.dumps({'tenantId': 'tenant_id', 'subscriptionId': 'subscription_id'}),
            )
        )

        self.resources = ResourceRequirements(requests=ResourceRequests(memory_in_gb='4', cpu='1'))
        with patch(
            'azure.common.credentials.ServicePrincipalCredentials.__init__', autospec=True, return_value=None
        ):
            with patch('azure.mgmt.containerinstance.ContainerInstanceManagementClient'):
                self.hook = AzureContainerInstanceHook(conn_id='azure_container_instance_test')

    @patch('azure.mgmt.containerinstance.models.ContainerGroup')
    @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.create_or_update')
    def test_create_or_update(self, create_or_update_mock, container_group_mock):
        self.hook.create_or_update('resource_group', 'aci-test', container_group_mock)
        create_or_update_mock.assert_called_once_with('resource_group', 'aci-test', container_group_mock)

    @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.get')
    def test_get_state(self, get_state_mock):
        self.hook.get_state('resource_group', 'aci-test')
        get_state_mock.assert_called_once_with('resource_group', 'aci-test', raw=False)

    @patch('azure.mgmt.containerinstance.operations.ContainerOperations.list_logs')
    def test_get_logs(self, list_logs_mock):
        expected_messages = ['log line 1\n', 'log line 2\n', 'log line 3\n']
        logs = Logs(content=''.join(expected_messages))
        list_logs_mock.return_value = logs

        logs = self.hook.get_logs('resource_group', 'name', 'name')

        assert logs == expected_messages

    @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.delete')
    def test_delete(self, delete_mock):
        self.hook.delete('resource_group', 'aci-test')
        delete_mock.assert_called_once_with('resource_group', 'aci-test')

    @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.list_by_resource_group')
    def test_exists_with_existing(self, list_mock):
        list_mock.return_value = [
            ContainerGroup(
                os_type='Linux',
                containers=[Container(name='test1', image='hello-world', resources=self.resources)],
            )
        ]
        assert not self.hook.exists('test', 'test1')

    @patch('azure.mgmt.containerinstance.operations.ContainerGroupsOperations.list_by_resource_group')
    def test_exists_with_not_existing(self, list_mock):
        list_mock.return_value = [
            ContainerGroup(
                os_type='Linux',
                containers=[Container(name='test1', image='hello-world', resources=self.resources)],
            )
        ]
        assert not self.hook.exists('test', 'not found')