Use unittest.mock instead of backported mock library (apache#11643)

mock is now part of the Python standard library, available as unittest.mock in Python 3.3 onwards.

tests/api_connexion/endpoints/test_config_endpoint.py

@@ -17,7 +17,7 @@
 
 import textwrap
 
-from mock import patch
+from unittest.mock import patch
 
 from airflow.security import permissions
 from airflow.www import app

tests/cli/commands/test_celery_command.py

@@ -18,8 +18,8 @@
 import unittest
 from argparse import Namespace
 from tempfile import NamedTemporaryFile
+from unittest import mock
 
-import mock
 import pytest
 import sqlalchemy
 
@@ -137,8 +137,9 @@ def test_same_pid_file_is_used_in_start_and_stop(
         celery_command.worker(worker_args)
         run_mock = mock_celery_worker.return_value.run
         assert run_mock.call_args
-        assert 'pidfile' in run_mock.call_args.kwargs
-        assert run_mock.call_args.kwargs['pidfile'] == pid_file
+        _, kwargs = run_mock.call_args
+        assert 'pidfile' in kwargs
+        assert kwargs['pidfile'] == pid_file
 
         # Call stop
         stop_args = self.parser.parse_args(['celery', 'stop'])

tests/cli/commands/test_dag_command.py

@@ -21,8 +21,7 @@
 import tempfile
 import unittest
 from datetime import datetime, timedelta
-
-import mock
+from unittest import mock
 
 from airflow import settings
 from airflow.cli import cli_parser

tests/core/test_logging_config.py

@@ -23,8 +23,7 @@
 import sys
 import tempfile
 import unittest
-
-from mock import patch
+from unittest.mock import patch
 
 from airflow.configuration import conf
 from tests.test_utils.config import conf_vars

tests/core/test_sqlalchemy_config.py

@@ -17,8 +17,8 @@
 # under the License.
 
 import unittest
+from unittest.mock import patch
 
-from mock import patch
 from sqlalchemy.pool import NullPool
 
 from airflow import settings

tests/executors/test_celery_kubernetes_executor.py

@@ -15,7 +15,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-import mock
+from unittest import mock
 
 from airflow.executors.celery_kubernetes_executor import CeleryKubernetesExecutor
 

tests/executors/test_executor_loader.py

@@ -16,8 +16,8 @@
 # under the License.
 
 import unittest
+from unittest import mock
 
-import mock
 from parameterized import parameterized
 
 from airflow import plugins_manager

tests/executors/test_kubernetes_executor.py

@@ -20,8 +20,8 @@
 import string
 import unittest
 from datetime import datetime
+from unittest import mock
 
-import mock
 from kubernetes.client import models as k8s
 from urllib3 import HTTPResponse
 

tests/jobs/test_backfill_job.py

@@ -22,10 +22,10 @@
 import logging
 import threading
 import unittest
+from unittest.mock import patch
 
 import pytest
 import sqlalchemy
-from mock import patch
 from parameterized import parameterized
 
 from airflow import settings

tests/jobs/test_base_job.py

@@ -18,8 +18,8 @@
 #
 
 import datetime
+from unittest.mock import ANY, Mock, patch
 
-from mock import ANY, Mock, patch
 from pytest import raises
 from sqlalchemy.exc import OperationalError
 

tests/jobs/test_local_task_job.py

@@ -22,9 +22,9 @@
 import unittest
 import uuid
 from unittest import mock
+from unittest.mock import patch
 
 import pytest
-from mock import patch
 
 from airflow import settings
 from airflow.exceptions import AirflowException

tests/jobs/test_scheduler_job.py

@@ -23,12 +23,12 @@
 import unittest
 from datetime import timedelta
 from tempfile import NamedTemporaryFile, mkdtemp
+from unittest import mock
+from unittest.mock import MagicMock, patch
 from zipfile import ZipFile
 
-import mock
 import psutil
 import pytest
-from mock import MagicMock, patch
 from parameterized import parameterized
 from sqlalchemy import func
 from sqlalchemy.exc import OperationalError

tests/kubernetes/test_client.py

@@ -17,8 +17,8 @@
 
 import socket
 import unittest
+from unittest import mock
 
-import mock
 from urllib3.connection import HTTPConnection, HTTPSConnection
 
 from airflow.kubernetes.kube_client import RefreshConfiguration, _enable_tcp_keepalive, get_kube_client

tests/kubernetes/test_pod_launcher.py

@@ -15,8 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 import unittest
+from unittest import mock
 
-import mock
 from requests.exceptions import BaseHTTPError
 
 from airflow.exceptions import AirflowException

tests/models/test_dagcode.py

@@ -17,8 +17,7 @@
 # under the License.
 import unittest
 from datetime import timedelta
-
-from mock import patch
+from unittest.mock import patch
 
 from airflow import AirflowException, example_dags as example_dags_module
 from airflow.models import DagBag

tests/models/test_dagrun.py

@@ -18,8 +18,8 @@
 
 import datetime
 import unittest
+from unittest import mock
 
-import mock
 from parameterized import parameterized
 
 from airflow import models, settings

tests/models/test_taskinstance.py

@@ -1492,7 +1492,7 @@ def on_execute_callable(context):
         assert ti.state == State.SUCCESS
 
     def test_handle_failure(self):
-        import mock
+        from unittest import mock
 
         start_date = timezone.datetime(2016, 6, 1)
         dag = models.DAG(dag_id="test_handle_failure", schedule_interval=None, start_date=start_date)

tests/operators/test_bash.py

@@ -17,12 +17,10 @@
 # under the License.
 
 import unittest
-import unittest.mock
 from datetime import datetime, timedelta
 from subprocess import PIPE, STDOUT
 from tempfile import NamedTemporaryFile
-
-import mock
+from unittest import mock
 
 from airflow.exceptions import AirflowException
 from airflow.models import DagRun
@@ -76,7 +74,7 @@ def test_echo_env_variables(self):
                              'echo $AIRFLOW_CTX_DAG_RUN_ID>> {0};'.format(tmp_file.name)
             )
 
-            with unittest.mock.patch.dict('os.environ', {
+            with mock.patch.dict('os.environ', {
                 'AIRFLOW_HOME': 'MY_PATH_TO_AIRFLOW_HOME',
                 'PYTHONPATH': 'AWESOME_PYTHONPATH'
             }):

tests/operators/test_sql.py

@@ -18,8 +18,8 @@
 
 import datetime
 import unittest
+from unittest import mock
 
-import mock
 import pytest
 
 from airflow.exceptions import AirflowException

tests/providers/amazon/aws/hooks/test_base_aws.py

@@ -17,9 +17,9 @@
 # under the License.
 #
 import unittest
+from unittest import mock
 
 import boto3
-import mock
 
 from airflow.models import Connection
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook

tests/providers/amazon/aws/hooks/test_batch_client.py

@@ -19,9 +19,9 @@
 # pylint: disable=missing-docstring
 
 import unittest
+from unittest import mock
 
 import botocore.exceptions
-import mock
 from parameterized import parameterized
 
 from airflow.exceptions import AirflowException

tests/providers/amazon/aws/hooks/test_batch_waiters.py

@@ -35,12 +35,12 @@
 import inspect
 import unittest
 from typing import NamedTuple, Optional
+from unittest import mock
 
 import boto3
 import botocore.client
 import botocore.exceptions
 import botocore.waiter
-import mock
 import pytest
 from moto import mock_batch, mock_ec2, mock_ecs, mock_iam, mock_logs
 

tests/providers/amazon/aws/hooks/test_datasync.py

@@ -17,9 +17,9 @@
 # under the License.
 #
 import unittest
+from unittest import mock
 
 import boto3
-import mock
 
 from airflow.exceptions import AirflowTaskTimeout
 from airflow.providers.amazon.aws.hooks.datasync import AWSDataSyncHook

tests/providers/amazon/aws/hooks/test_glacier.py

@@ -18,7 +18,7 @@
 
 import unittest
 
-import mock
+from unittest import mock
 from testfixtures import LogCapture
 
 from airflow.providers.amazon.aws.hooks.glacier import GlacierHook

tests/providers/amazon/aws/hooks/test_glue.py

@@ -18,7 +18,7 @@
 import json
 import unittest
 
-import mock
+from unittest import mock
 
 from airflow.providers.amazon.aws.hooks.glue import AwsGlueJobHook
 

tests/providers/amazon/aws/hooks/test_glue_catalog.py

@@ -17,9 +17,9 @@
 # under the License.
 
 import unittest
+from unittest import mock
 
 import boto3
-import mock
 
 from airflow.providers.amazon.aws.hooks.glue_catalog import AwsGlueCatalogHook
 

tests/providers/amazon/aws/hooks/test_s3.py

@@ -19,10 +19,10 @@
 import gzip as gz
 import os
 import tempfile
+from unittest import mock
 from unittest.mock import Mock
 
 import boto3
-import mock
 import pytest
 from botocore.exceptions import ClientError, NoCredentialsError
 

tests/providers/amazon/aws/hooks/test_sagemaker.py

@@ -21,7 +21,7 @@
 import unittest
 from datetime import datetime
 
-import mock
+from unittest import mock
 from tzlocal import get_localzone
 
 from airflow.exceptions import AirflowException

tests/providers/amazon/aws/operators/test_athena.py

@@ -17,7 +17,7 @@
 
 import unittest
 
-import mock
+from unittest import mock
 
 from airflow.models import DAG, TaskInstance
 from airflow.providers.amazon.aws.hooks.athena import AWSAthenaHook

tests/providers/amazon/aws/operators/test_batch.py

@@ -21,7 +21,7 @@
 
 import unittest
 
-import mock
+from unittest import mock
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.batch_client import AwsBatchClientHook

tests/providers/amazon/aws/operators/test_ecs.py

@@ -21,7 +21,7 @@
 import unittest
 from copy import deepcopy
 
-import mock
+from unittest import mock
 from parameterized import parameterized
 
 from airflow.exceptions import AirflowException

tests/providers/amazon/aws/operators/test_glacier.py

@@ -18,7 +18,7 @@
 
 from unittest import TestCase
 
-import mock
+from unittest import mock
 
 from airflow.providers.amazon.aws.operators.glacier import (
     GlacierCreateJobOperator,

tests/providers/amazon/aws/operators/test_glue.py

@@ -17,7 +17,7 @@
 
 import unittest
 
-import mock
+from unittest import mock
 
 from airflow import configuration
 from airflow.providers.amazon.aws.hooks.glue import AwsGlueJobHook

tests/providers/amazon/aws/operators/test_s3_bucket.py

@@ -18,7 +18,7 @@
 import os
 import unittest
 
-import mock
+from unittest import mock
 from moto import mock_s3
 
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook

tests/providers/amazon/aws/operators/test_s3_list.py

@@ -18,7 +18,7 @@
 
 import unittest
 
-import mock
+from unittest import mock
 
 from airflow.providers.amazon.aws.operators.s3_list import S3ListOperator
 

tests/providers/amazon/aws/operators/test_sagemaker_endpoint.py

@@ -18,7 +18,7 @@
 
 import unittest
 
-import mock
+from unittest import mock
 from botocore.exceptions import ClientError
 
 from airflow.exceptions import AirflowException

tests/providers/amazon/aws/operators/test_sagemaker_endpoint_config.py

@@ -18,7 +18,7 @@
 
 import unittest
 
-import mock
+from unittest import mock
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook