[AIRFLOW-4215] Replace mock with unittest.mock (#5292)

setup.py

@@ -260,7 +260,6 @@ devel = [
     'flake8>=3.6.0',
     'freezegun',
     'jira',
-    'mock;python_version<"3.3"',
     'mongomock',
     'moto==1.3.5',
     'nose',

tests/api/client/test_local_client.py

@@ -21,7 +21,7 @@ import json
 import unittest
 
 from freezegun import freeze_time
-from mock import patch
+from unittest.mock import patch
 
 from airflow import AirflowException
 from airflow import models

tests/api/common/experimental/test_trigger_dag.py

@@ -17,7 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import mock
+from unittest import mock
 import unittest
 import json
 

tests/cli/test_cli.py

@@ -23,7 +23,7 @@ import sys
 import unittest
 
 from datetime import datetime, timedelta, time
-from mock import patch, Mock, MagicMock
+from unittest.mock import patch, Mock, MagicMock
 from time import sleep
 import psutil
 import pytz

tests/contrib/hooks/test_azure_container_instance_hook.py

@@ -20,7 +20,7 @@
 import json
 import unittest
 from collections import namedtuple
-from mock import patch
+from unittest.mock import patch
 
 from airflow import configuration
 from airflow.models import Connection

tests/contrib/hooks/test_bigquery_hook.py

@@ -22,7 +22,7 @@ import unittest
 from typing import List
 
 from google.auth.exceptions import GoogleAuthError
-import mock
+from unittest import mock
 from googleapiclient.errors import HttpError
 
 from airflow.contrib.hooks import bigquery_hook as hook

tests/contrib/hooks/test_datadog_hook.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 import json
-import mock
+from unittest import mock
 import unittest
 
 from airflow.exceptions import AirflowException

tests/contrib/hooks/test_datastore_hook.py

@@ -21,7 +21,7 @@
 import unittest
 
 from airflow.contrib.hooks.datastore_hook import DatastoreHook
-from mock import call, patch
+from unittest.mock import call, patch
 
 from tests.compat import mock
 

tests/contrib/hooks/test_ftp_hook.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 
-import mock
+from unittest import mock
 import six
 import unittest
 

tests/contrib/hooks/test_gcp_mlengine_hook.py

@@ -18,7 +18,7 @@
 import json
 import unittest
 
-import mock
+from unittest import mock
 import requests
 from google.auth.exceptions import GoogleAuthError
 from googleapiclient.discovery import build_from_document

tests/contrib/hooks/test_imap_hook.py

@@ -20,7 +20,7 @@
 import imaplib
 import unittest
 
-from mock import Mock, patch, mock_open
+from unittest.mock import Mock, patch, mock_open
 
 from airflow import configuration, AirflowException
 from airflow.contrib.hooks.imap_hook import ImapHook

tests/contrib/hooks/test_jdbc_hook.py

@@ -21,8 +21,8 @@
 import unittest
 import json
 
-from mock import Mock
+from unittest.mock import Mock
-from mock import patch
+from unittest.mock import patch
 
 from airflow import configuration
 from airflow.hooks.jdbc_hook import JdbcHook

tests/contrib/hooks/test_jira_hook.py

@@ -20,8 +20,8 @@
 
 import unittest
 
-from mock import Mock
+from unittest.mock import Mock
-from mock import patch
+from unittest.mock import patch
 
 from airflow import configuration
 from airflow.contrib.hooks.jira_hook import JiraHook

tests/contrib/hooks/test_pinot_hook.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 
-import mock
+from unittest import mock
 import unittest
 
 from airflow.contrib.hooks.pinot_hook import PinotDbApiHook

tests/contrib/hooks/test_salesforce_hook.py

@@ -21,7 +21,7 @@
 import unittest
 
 import pandas as pd
-from mock import patch, Mock
+from unittest.mock import patch, Mock
 from simple_salesforce import Salesforce
 
 from airflow.contrib.hooks.salesforce_hook import SalesforceHook

tests/contrib/hooks/test_segment_hook.py

@@ -17,7 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 #
-import mock
+from unittest import mock
 import unittest
 
 from airflow import configuration, AirflowException

tests/contrib/hooks/test_sftp_hook.py

@@ -17,7 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import mock
+from unittest import mock
 import unittest
 import shutil
 import os

tests/contrib/hooks/test_slack_webhook_hook.py

@@ -20,12 +20,12 @@
 import json
 from requests.exceptions import MissingSchema
 import unittest
+from unittest import mock
 
 from airflow import configuration
 from airflow.models import Connection
 from airflow.utils import db
 from airflow.contrib.hooks.slack_webhook_hook import SlackWebhookHook
-from tests.compat import mock
 
 
 class TestSlackWebhookHook(unittest.TestCase):
@@ -104,8 +104,9 @@ class TestSlackWebhookHook(unittest.TestCase):
         # Then
         self.assertEqual(self.expected_message, message)
 
+    @mock.patch('requests.Session')
     @mock.patch('requests.Request')
-    def test_url_generated_by_http_conn_id(self, request_mock):
+    def test_url_generated_by_http_conn_id(self, request_mock, session_mock):
         hook = SlackWebhookHook(http_conn_id='slack-webhook-url')
         try:
             hook.execute()
@@ -119,8 +120,9 @@ class TestSlackWebhookHook(unittest.TestCase):
         )
         request_mock.reset_mock()
 
+    @mock.patch('requests.Session')
     @mock.patch('requests.Request')
-    def test_url_generated_by_endpoint(self, request_mock):
+    def test_url_generated_by_endpoint(self, request_mock, session_mock):
         hook = SlackWebhookHook(webhook_token=self.expected_url)
         try:
             hook.execute()
@@ -134,8 +136,9 @@ class TestSlackWebhookHook(unittest.TestCase):
         )
         request_mock.reset_mock()
 
+    @mock.patch('requests.Session')
     @mock.patch('requests.Request')
-    def test_url_generated_by_http_conn_id_and_endpoint(self, request_mock):
+    def test_url_generated_by_http_conn_id_and_endpoint(self, request_mock, session_mock):
         hook = SlackWebhookHook(http_conn_id='slack-webhook-host',
                                 webhook_token='B000/XXX')
         try:

tests/contrib/hooks/test_snowflake_hook.py

@@ -19,7 +19,7 @@
 #
 import os
 
-import mock
+from unittest import mock
 import unittest
 
 from cryptography.hazmat.primitives import serialization

tests/contrib/hooks/test_spark_sql_hook.py

@@ -22,7 +22,7 @@ import six
 import unittest
 from itertools import dropwhile
 
-from mock import patch, call
+from unittest.mock import patch, call
 
 from airflow import configuration
 from airflow.models import Connection

tests/contrib/hooks/test_spark_submit_hook.py

@@ -23,7 +23,7 @@ import unittest
 from airflow import configuration, AirflowException
 from airflow.models import Connection
 from airflow.utils import db
-from mock import patch, call
+from unittest.mock import patch, call
 
 from airflow.contrib.hooks.spark_submit_hook import SparkSubmitHook
 

tests/contrib/hooks/test_sqoop_hook.py

@@ -28,7 +28,7 @@ from airflow.exceptions import AirflowException
 from airflow.models import Connection
 from airflow.utils import db
 
-from mock import patch, call
+from unittest.mock import patch, call
 
 from io import StringIO
 

tests/contrib/hooks/test_winrm_hook.py

@@ -20,7 +20,7 @@
 
 import unittest
 
-from mock import patch
+from unittest.mock import patch
 
 from airflow import AirflowException
 from airflow.contrib.hooks.winrm_hook import WinRMHook

tests/contrib/hooks/test_zendesk_hook.py

@@ -20,7 +20,7 @@
 
 import unittest
 
-import mock
+from unittest import mock
 
 from airflow.hooks.zendesk_hook import ZendeskHook
 from zdesk import RateLimitError

tests/contrib/operators/test_aws_sqs_publish_operator.py

@@ -22,7 +22,7 @@ import unittest
 from airflow import DAG, configuration
 from airflow.contrib.operators.aws_sqs_publish_operator import SQSPublishOperator
 from airflow.utils import timezone
-from mock import MagicMock
+from unittest.mock import MagicMock
 from moto import mock_sqs
 from airflow.contrib.hooks.aws_sqs_hook import SQSHook
 

tests/contrib/operators/test_cassandra_to_gcs_operator.py

@@ -18,7 +18,7 @@
 # under the License.
 
 import unittest
-import mock
+from unittest import mock
 from builtins import str
 from airflow.contrib.operators.cassandra_to_gcs import \
     CassandraToGoogleCloudStorageOperator

tests/contrib/operators/test_dataproc_operator.py

@@ -39,8 +39,8 @@ from tests.compat import mock
 
 from copy import deepcopy
 
-from mock import MagicMock, Mock
+from unittest.mock import MagicMock, Mock
-from mock import patch
+from unittest.mock import patch
 
 TASK_ID = 'test-dataproc-operator'
 CLUSTER_NAME = 'test-cluster-name'

tests/contrib/operators/test_dingding_operator.py

@@ -19,7 +19,7 @@
 
 import unittest
 
-import mock
+from unittest import mock
 
 from airflow import DAG, configuration
 from airflow.contrib.operators.dingding_operator import DingdingOperator
@@ -68,4 +68,4 @@ class TestDingdingOperator(unittest.TestCase):
             self._config['at_mobiles'],
             self._config['at_all']
         )
-        mock_hook.return_value.send.assert_called_once()
+        mock_hook.return_value.send.assert_called_once_with()

tests/contrib/operators/test_druid_operator.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 
-import mock
+from unittest import mock
 import unittest
 
 from airflow import DAG, configuration

tests/contrib/operators/test_emr_add_steps_operator.py

@@ -20,7 +20,7 @@
 import unittest
 from datetime import timedelta
 
-from mock import MagicMock, patch
+from unittest.mock import MagicMock, patch
 
 from airflow import DAG, configuration
 from airflow.contrib.operators.emr_add_steps_operator import EmrAddStepsOperator

tests/contrib/operators/test_emr_create_job_flow_operator.py

@@ -21,7 +21,7 @@
 import unittest
 from datetime import timedelta
 
-from mock import MagicMock, patch
+from unittest.mock import MagicMock, patch
 
 from airflow import DAG, configuration
 from airflow.contrib.operators.emr_create_job_flow_operator import EmrCreateJobFlowOperator

tests/contrib/operators/test_emr_terminate_job_flow_operator.py

@@ -18,7 +18,7 @@
 # under the License.
 
 import unittest
-from mock import MagicMock, patch
+from unittest.mock import MagicMock, patch
 
 from airflow import configuration
 from airflow.contrib.operators.emr_terminate_job_flow_operator import EmrTerminateJobFlowOperator

tests/contrib/operators/test_hive_to_dynamodb_operator.py

@@ -22,7 +22,7 @@ import json
 import unittest
 import datetime
 
-import mock
+from unittest import mock
 import pandas as pd
 
 from airflow import configuration, DAG

tests/contrib/operators/test_imap_attachment_to_s3_operator.py

@@ -19,7 +19,7 @@
 
 import unittest
 
-from mock import patch
+from unittest.mock import patch
 
 from airflow.contrib.operators.imap_attachment_to_s3_operator import ImapAttachmentToS3Operator
 

tests/contrib/operators/test_jira_operator_test.py

@@ -20,8 +20,8 @@
 
 import unittest
 
-from mock import Mock
+from unittest.mock import Mock
-from mock import patch
+from unittest.mock import patch
 
 from airflow import DAG, configuration
 from airflow.contrib.operators.jira_operator import JiraOperator

tests/contrib/operators/test_mlengine_operator.py

@@ -21,7 +21,7 @@ import unittest
 
 import httplib2
 from googleapiclient.errors import HttpError
-from mock import ANY, patch
+from unittest.mock import ANY, patch
 
 from airflow import DAG, configuration
 from airflow.contrib.operators.mlengine_operator import (MLEngineBatchPredictionOperator,

tests/contrib/operators/test_segment_track_event_operator.py

@@ -17,7 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 #
-import mock
+from unittest import mock
 import unittest
 
 from airflow import configuration, AirflowException

tests/contrib/operators/test_sns_publish_operator.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 
-import mock
+from unittest import mock
 import unittest
 
 from airflow.contrib.operators.sns_publish_operator import SnsPublishOperator

tests/contrib/operators/test_vertica_operator.py

@@ -17,7 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import mock
+from unittest import mock
 import unittest
 
 from airflow.contrib.operators.vertica_operator import VerticaOperator

tests/contrib/operators/test_vertica_to_mysql.py

@@ -19,7 +19,7 @@
 
 import datetime
 
-import mock
+from unittest import mock
 import unittest
 
 from airflow import DAG, configuration

tests/contrib/operators/test_winrm_operator.py

@@ -17,7 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import mock
+from unittest import mock
 import unittest
 
 from airflow.contrib.operators.winrm_operator import WinRMOperator

tests/contrib/sensors/test_cassandra_sensor.py

@@ -20,7 +20,7 @@
 
 import unittest
 
-from mock import patch
+from unittest.mock import patch
 
 from airflow import DAG
 from airflow import configuration

tests/contrib/sensors/test_celery_queue_sensor.py

@@ -18,7 +18,7 @@
 # under the License.
 
 import unittest
-from mock import patch
+from unittest.mock import patch
 
 from airflow.contrib.sensors.celery_queue_sensor import CeleryQueueSensor
 

tests/contrib/sensors/test_emr_job_flow_sensor.py

@@ -20,7 +20,7 @@
 import unittest
 import datetime
 from dateutil.tz import tzlocal
-from mock import MagicMock, patch
+from unittest.mock import MagicMock, patch
 
 from airflow import configuration
 from airflow.contrib.sensors.emr_job_flow_sensor import EmrJobFlowSensor

tests/contrib/sensors/test_emr_step_sensor.py

@@ -19,7 +19,7 @@
 
 import unittest
 from datetime import datetime
-from mock import MagicMock, patch
+from unittest.mock import MagicMock, patch
 from dateutil.tz import tzlocal
 
 from airflow import configuration, AirflowException

tests/contrib/sensors/test_ftp_sensor.py

@@ -20,7 +20,7 @@
 import unittest
 
 from ftplib import error_perm
-from mock import MagicMock
+from unittest.mock import MagicMock
 
 from airflow.contrib.hooks.ftp_hook import FTPHook
 from airflow.contrib.sensors.ftp_sensor import FTPSensor

tests/contrib/sensors/test_imap_attachment_sensor.py

@@ -19,7 +19,7 @@
 
 import unittest
 
-from mock import patch, Mock
+from unittest.mock import patch, Mock
 
 from airflow import configuration
 from airflow.contrib.sensors.imap_attachment_sensor import ImapAttachmentSensor

tests/contrib/sensors/test_jira_sensor_test.py

@@ -20,8 +20,8 @@
 
 import unittest
 
-from mock import Mock
+from unittest.mock import Mock
-from mock import patch
+from unittest.mock import patch
 
 from airflow import DAG, configuration
 from airflow.contrib.sensors.jira_sensor import JiraTicketSensor

tests/contrib/sensors/test_qubole_sensor.py

@@ -21,7 +21,7 @@
 import unittest
 
 from datetime import datetime
-from mock import patch
+from unittest.mock import patch
 
 from airflow.contrib.sensors.qubole_sensor import QuboleFileSensor, QubolePartitionSensor
 from airflow.exceptions import AirflowException

tests/contrib/sensors/test_redis_pub_sub_sensor.py

@@ -23,7 +23,7 @@ from airflow import DAG, configuration
 from airflow.contrib.sensors.redis_pub_sub_sensor import RedisPubSubSensor
 from airflow.utils import timezone
 from airflow.contrib.hooks.redis_hook import RedisHook
-from mock import patch, call, MagicMock
+from unittest.mock import patch, call, MagicMock
 
 DEFAULT_DATE = timezone.datetime(2017, 1, 1)
 

tests/contrib/sensors/test_sftp_sensor.py

@@ -18,7 +18,7 @@
 # under the License.
 
 import unittest
-from mock import patch
+from unittest.mock import patch
 from paramiko import SFTP_NO_SUCH_FILE, SFTP_FAILURE
 from airflow.contrib.sensors.sftp_sensor import SFTPSensor
 

tests/contrib/sensors/test_sqs_sensor.py

@@ -22,7 +22,7 @@ import unittest
 from airflow import DAG, configuration
 from airflow.contrib.sensors.aws_sqs_sensor import SQSSensor
 from airflow.utils import timezone
-from mock import patch, MagicMock
+from unittest.mock import patch, MagicMock
 from airflow.exceptions import AirflowException
 from moto import mock_sqs
 from airflow.contrib.hooks.aws_sqs_hook import SQSHook

tests/contrib/task_runner/test_cgroup_task_runner.py

@@ -16,7 +16,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
 import unittest
 
 from airflow.contrib.task_runner.cgroup_task_runner import CgroupTaskRunner

tests/contrib/utils/base_gcp_mock.py

@@ -17,7 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import mock
+from unittest import mock
 
 GCP_PROJECT_ID_HOOK_UNIT_TEST = 'example-project'
 

tests/contrib/utils/test_mlengine_operator_utils.py

@@ -23,8 +23,8 @@ from airflow.contrib.utils import mlengine_operator_utils
 from airflow.exceptions import AirflowException
 from airflow.version import version
 
-from mock import ANY
+from unittest.mock import ANY
-from mock import patch
+from unittest.mock import patch
 
 DEFAULT_DATE = datetime.datetime(2017, 6, 6)
 TEST_VERSION = 'v{}'.format(version.replace('.', '-').replace('+', '-'))

tests/core.py

@@ -21,7 +21,7 @@ import json
 import unittest
 
 import doctest
-import mock
+from unittest import mock
 import multiprocessing
 import os
 import pickle  # type: ignore

tests/executors/test_celery_executor.py

@@ -22,7 +22,7 @@ import unittest
 import contextlib
 from multiprocessing import Pool
 
-import mock
+from unittest import mock
 
 from celery import Celery
 from celery import states as celery_states

tests/hooks/test_dbapi_hook.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 
-import mock
+from unittest import mock
 import unittest
 
 from airflow.hooks.dbapi_hook import DbApiHook

tests/hooks/test_druid_hook.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 
-from mock import MagicMock, patch
+from unittest.mock import MagicMock, patch
 import requests
 import requests_mock
 import unittest

tests/hooks/test_hive_hook.py

@@ -26,7 +26,7 @@ import sys
 import unittest
 from collections import OrderedDict
 
-import mock
+from unittest import mock
 import pandas as pd
 from hmsclient import HMSClient
 

tests/hooks/test_http_hook.py

@@ -76,8 +76,9 @@ class TestHttpHook(unittest.TestCase):
             self.assertEqual(resp.text, '{"status":{"status": 200}}')
 
     @requests_mock.mock()
+    @mock.patch('requests.Session')
     @mock.patch('requests.Request')
-    def test_get_request_with_port(self, m, request_mock):
+    def test_get_request_with_port(self, m, request_mock, session_mock):
         from requests.exceptions import MissingSchema
 
         with mock.patch(

tests/hooks/test_mysql_hook.py

@@ -19,7 +19,7 @@
 #
 
 import json
-import mock
+from unittest import mock
 import unittest
 
 import MySQLdb.cursors

tests/hooks/test_postgres_hook.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 
-import mock
+from unittest import mock
 import unittest
 
 from tempfile import NamedTemporaryFile

tests/hooks/test_presto_hook.py

@@ -18,10 +18,10 @@
 # under the License.
 #
 
-import mock
+from unittest import mock
 import unittest
 
-from mock import patch
+from unittest.mock import patch
 
 from airflow.hooks.presto_hook import PrestoHook
 

tests/hooks/test_s3_hook.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 
-import mock
+from unittest import mock
 import tempfile
 import unittest
 

tests/hooks/test_webhdfs_hook.py

@@ -20,7 +20,7 @@
 import unittest
 
 from hdfs import HdfsError
-from mock import patch, call
+from unittest.mock import patch, call
 
 from airflow.hooks.webhdfs_hook import WebHDFSHook, AirflowWebHDFSHookException
 from airflow.models.connection import Connection

tests/models/test_connection.py

@@ -21,7 +21,7 @@ import unittest
 from collections import namedtuple
 
 from cryptography.fernet import Fernet
-from mock import patch
+from unittest.mock import patch
 from parameterized import parameterized
 
 from airflow.models import Connection, crypto

tests/models/test_dag.py

@@ -28,7 +28,7 @@ from tempfile import NamedTemporaryFile
 import jinja2
 import pendulum
 import six
-from mock import patch
+from unittest.mock import patch
 
 from airflow import models, settings, configuration
 from airflow.exceptions import AirflowException, AirflowDagCycleException

tests/models/test_dagbag.py

@@ -24,7 +24,7 @@ import textwrap
 import unittest
 from tempfile import mkdtemp, NamedTemporaryFile
 
-from mock import patch, ANY
+from unittest.mock import patch, ANY
 
 from airflow import models, configuration
 from airflow.models import DagModel, DagBag, TaskInstance as TI
@@ -128,7 +128,7 @@ class DagBagTest(unittest.TestCase):
         test the loading of a DAG from within a zip file that skips another file because
         it doesn't have "airflow" and "DAG"
         """
-        from mock import Mock
+        from unittest.mock import Mock
         with patch('airflow.models.DagBag.log') as log_mock:
             log_mock.info = Mock()
             test_zip_path = os.path.join(TEST_DAGS_FOLDER, "test_zip.zip")

tests/models/test_kubernetes.py

@@ -19,7 +19,7 @@
 
 import unittest
 
-from mock import patch
+from unittest.mock import patch
 
 from airflow import settings
 from airflow.models import KubeResourceVersion, KubeWorkerIdentifier

tests/models/test_skipmixin.py

@@ -21,7 +21,7 @@ import datetime
 import unittest
 
 import pendulum
-from mock import patch, Mock
+from unittest.mock import patch, Mock
 
 from airflow import settings
 from airflow.models import DAG, TaskInstance as TI

tests/models/test_taskinstance.py

@@ -24,7 +24,7 @@ import urllib
 
 import pendulum
 from freezegun import freeze_time
-from mock import patch, mock_open
+from unittest.mock import patch, mock_open
 from parameterized import parameterized
 
 from airflow import models, settings, configuration

tests/models/test_variable.py

@@ -20,7 +20,7 @@
 import unittest
 
 from cryptography.fernet import Fernet
-from mock import patch
+from unittest.mock import patch
 
 from airflow import settings
 from airflow.models import crypto, Variable

tests/operators/test_email_operator.py

@@ -18,7 +18,7 @@
 # under the License.
 
 import datetime
-import mock
+from unittest import mock
 import unittest
 
 from airflow import configuration, DAG

tests/operators/test_hive_operator.py

@@ -21,7 +21,7 @@ import datetime
 import os
 import unittest
 
-import mock
+from unittest import mock
 import nose
 
 from airflow import DAG, configuration, operators

tests/operators/test_hive_stats_operator.py

@@ -20,7 +20,7 @@
 import unittest
 from collections import OrderedDict
 
-from mock import patch
+from unittest.mock import patch
 
 from airflow import AirflowException
 from airflow.operators.hive_stats_operator import HiveStatsCollectionOperator

tests/operators/test_hive_to_mysql.py

@@ -19,7 +19,7 @@
 
 import unittest
 
-from mock import patch, PropertyMock
+from unittest.mock import patch, PropertyMock
 
 from airflow.operators.hive_to_mysql import HiveToMySqlTransfer
 from airflow.utils.operator_helpers import context_to_airflow_vars

tests/operators/test_hive_to_samba_operator.py

@@ -19,7 +19,7 @@
 
 import unittest
 
-from mock import patch, PropertyMock, Mock
+from unittest.mock import patch, PropertyMock, Mock
 
 from airflow.operators.hive_to_samba_operator import Hive2SambaOperator
 from airflow.utils.operator_helpers import context_to_airflow_vars

tests/operators/test_jdbc_operator.py

@@ -19,7 +19,7 @@
 
 import unittest
 
-from mock import patch
+from unittest.mock import patch
 
 from airflow.operators.jdbc_operator import JdbcOperator
 

tests/operators/test_mssql_to_hive.py

@@ -24,7 +24,7 @@ try:
     import pymssql
 except ImportError:
     pymssql = None
-from mock import patch, PropertyMock, Mock
+from unittest.mock import patch, PropertyMock, Mock
 
 from airflow.operators.mssql_to_hive import MsSqlToHiveTransfer
 

tests/operators/test_operators.py

@@ -23,7 +23,7 @@ from airflow.utils import timezone
 from collections import OrderedDict
 
 import os
-import mock
+from unittest import mock
 import unittest
 
 configuration.load_test_config()

tests/operators/test_papermill_operator.py

@@ -18,7 +18,7 @@
 # under the License.
 import unittest
 
-from mock import patch
+from unittest.mock import patch
 
 from airflow.operators.papermill_operator import PapermillOperator
 

tests/operators/test_presto_to_mysql.py

@@ -19,7 +19,7 @@
 
 import unittest
 
-from mock import patch
+from unittest.mock import patch
 
 from airflow.operators.presto_to_mysql import PrestoToMySqlTransfer
 

tests/operators/test_redshift_to_s3_operator.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 
-import mock
+from unittest import mock
 import unittest
 
 from boto3.session import Session

tests/operators/test_s3_file_transform_operator.py

@@ -27,7 +27,7 @@ import unittest
 from tempfile import mkdtemp
 
 import boto3
-import mock
+from unittest import mock
 from moto import mock_s3
 
 from airflow.exceptions import AirflowException

tests/operators/test_s3_to_redshift_operator.py

@@ -18,7 +18,7 @@
 # under the License.
 #
 
-import mock
+from unittest import mock
 import unittest
 
 from boto3.session import Session

tests/operators/test_subdag_operator.py

@@ -19,7 +19,7 @@
 
 import unittest
 
-from mock import Mock
+from unittest.mock import Mock
 
 import airflow
 from airflow.exceptions import AirflowException

tests/sensors/test_base_sensor.py

@@ -18,7 +18,7 @@
 # under the License.
 
 import unittest
-from mock import Mock
+from unittest.mock import Mock
 
 from airflow import DAG, configuration, settings
 from airflow.exceptions import (AirflowSensorTimeout, AirflowException,

tests/sensors/test_http_sensor.py

@@ -19,7 +19,7 @@
 import unittest
 
 import requests
-from mock import patch
+from unittest.mock import patch
 
 from airflow import DAG, configuration
 from airflow.exceptions import AirflowException, AirflowSensorTimeout

tests/sensors/test_s3_key_sensor.py

@@ -17,7 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import mock
+from unittest import mock
 import unittest
 from parameterized import parameterized
 

tests/sensors/test_s3_prefix_sensor.py

@@ -17,7 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import mock
+from unittest import mock
 import unittest
 
 from airflow.sensors.s3_prefix_sensor import S3PrefixSensor

tests/sensors/test_sql_sensor.py

@@ -16,7 +16,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
 import unittest
 
 from airflow import DAG

tests/task/task_runner/test_standard_task_runner.py

@@ -16,7 +16,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
 import os
 import psutil
 import time

tests/test_stats.py

@@ -20,7 +20,7 @@
 import unittest
 
 from airflow.stats import SafeStatsdLogger
-from mock import Mock
+from unittest.mock import Mock
 
 
 class TestStats(unittest.TestCase):

tests/test_utils/decorators.py

@@ -19,7 +19,7 @@
 
 from contextlib import ContextDecorator
 
-from mock import Mock
+from unittest.mock import Mock
 
 from airflow import conf
 

tests/ti_deps/deps/test_dag_ti_slots_available_dep.py

@@ -18,7 +18,7 @@
 # under the License.
 
 import unittest
-from mock import Mock
+from unittest.mock import Mock
 
 from airflow.models import TaskInstance
 from airflow.ti_deps.deps.dag_ti_slots_available_dep import DagTISlotsAvailableDep

tests/ti_deps/deps/test_dag_unpaused_dep.py

@@ -18,7 +18,7 @@
 # under the License.
 
 import unittest
-from mock import Mock
+from unittest.mock import Mock
 
 from airflow.models import TaskInstance
 from airflow.ti_deps.deps.dag_unpaused_dep import DagUnpausedDep

tests/ti_deps/deps/test_dagrun_exists_dep.py

@@ -19,7 +19,7 @@
 
 import unittest
 from airflow.utils.state import State
-from mock import Mock, patch
+from unittest.mock import Mock, patch
 
 from airflow.models import DAG, DagRun
 from airflow.ti_deps.deps.dagrun_exists_dep import DagrunRunningDep

tests/ti_deps/deps/test_not_in_retry_period_dep.py

@@ -20,7 +20,7 @@
 import unittest
 from datetime import timedelta
 from freezegun import freeze_time
-from mock import Mock
+from unittest.mock import Mock
 
 from airflow.models import TaskInstance
 from airflow.ti_deps.deps.not_in_retry_period_dep import NotInRetryPeriodDep

tests/ti_deps/deps/test_not_running_dep.py

@@ -19,7 +19,7 @@
 
 import unittest
 from datetime import datetime
-from mock import Mock
+from unittest.mock import Mock
 
 from airflow.ti_deps.deps.not_running_dep import NotRunningDep
 from airflow.utils.state import State