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

setup.py

@@ -260,7 +260,6 @@ def write_version(filename=os.path.join(*['airflow',
     '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 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 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 @@
 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 mock import patch
+from unittest.mock import Mock
+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 mock import patch
+from unittest.mock import Mock
+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 @@ def test_build_slack_message(self):
         # 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 @@ def test_url_generated_by_http_conn_id(self, request_mock):
         )
         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 @@ def test_url_generated_by_endpoint(self, request_mock):
         )
         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 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 @@
 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.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 @@
 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 copy import deepcopy
 
-from mock import MagicMock, Mock
-from mock import patch
+from unittest.mock import MagicMock, Mock
+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 @@ def test_execute(self, mock_hook):
             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 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 mock import patch
+from unittest.mock import Mock
+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 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