Fix DockerOperator & some operator test (apache#4049)

- For argument `image`, no need to explicitly
  add "latest" if tag is omitted.
  "latest" will be used by default if no
  tag provided. This is handled by `docker` package itself.

- Intermediate variable `cpu_shares` is not needed.

- Fix wrong usage of `cpu_shares` and `cpu_shares`.
  Based on
  https://docker-py.readthedocs.io/en/stable/api.html#docker.api.container.ContainerApiMixin.create_host_config,
  They should be an arguments of
  self.cli.create_host_config()
  rather than
  APIClient.create_container().

- Change name of the corresponding test script,
  to ensure it can be discovered.

- Fix the test itself.

- Some other test scripts are not named properly,
  which result in failure of test discovery.

Author: XD-DENG

airflow/operators/docker_operator.py

@@ -43,6 +43,7 @@ class DockerOperator(BaseOperator):
     be provided with the parameter ``docker_conn_id``.
 
     :param image: Docker image from which to create the container.
+        If image tag is omitted, "latest" will be used.
     :type image: str
     :param api_version: Remote API version. Set to ``auto`` to automatically
         detect the server's version.
@@ -62,7 +63,7 @@ class DockerOperator(BaseOperator):
     :type docker_url: str
     :param environment: Environment variables to set in the container. (templated)
     :type environment: dict
-    :param force_pull: Pull the docker image on every run. Default is false.
+    :param force_pull: Pull the docker image on every run. Default is False.
     :type force_pull: bool
     :param mem_limit: Maximum amount of memory the container can use.
         Either a float value, which represents the limit in bytes,
@@ -187,35 +188,28 @@ def execute(self, context):
                 tls=tls_config
             )
 
-        if ':' not in self.image:
-            image = self.image + ':latest'
-        else:
-            image = self.image
-
-        if self.force_pull or len(self.cli.images(name=image)) == 0:
-            self.log.info('Pulling docker image %s', image)
-            for l in self.cli.pull(image, stream=True):
+        if self.force_pull or len(self.cli.images(name=self.image)) == 0:
+            self.log.info('Pulling docker image %s', self.image)
+            for l in self.cli.pull(self.image, stream=True):
                 output = json.loads(l.decode('utf-8'))
                 self.log.info("%s", output['status'])
 
-        cpu_shares = int(round(self.cpus * 1024))
-
         with TemporaryDirectory(prefix='airflowtmp') as host_tmp_dir:
             self.environment['AIRFLOW_TMP_DIR'] = self.tmp_dir
             self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir))
 
             self.container = self.cli.create_container(
                 command=self.get_command(),
-                cpu_shares=cpu_shares,
                 environment=self.environment,
                 host_config=self.cli.create_host_config(
                     binds=self.volumes,
                     network_mode=self.network_mode,
                     shm_size=self.shm_size,
                     dns=self.dns,
-                    dns_search=self.dns_search),
-                image=image,
-                mem_limit=self.mem_limit,
+                    dns_search=self.dns_search,
+                    cpu_shares=int(round(self.cpus * 1024)),
+                    mem_limit=self.mem_limit),
+                image=self.image,
                 user=self.user,
                 working_dir=self.working_dir
             )

tests/operators/__init__.py

@@ -7,21 +7,12 @@
 # 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.
-
-from .docker_operator import *
-from .subdag_operator import *
-from .operators import *
-from .hive_operator import *
-from .s3_to_hive_operator import *
-from .python_operator import *
-from .latest_only_operator import *
-

tests/operators/docker_operator.py tests/operators/test_docker_operator.py

@@ -7,9 +7,9 @@
 # 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
@@ -64,20 +64,22 @@ def test_execute(self, client_class_mock, mkdtemp_mock):
         client_class_mock.assert_called_with(base_url='unix://var/run/docker.sock', tls=None,
                                              version='1.19')
 
-        client_mock.create_container.assert_called_with(command='env', cpu_shares=1024,
+        client_mock.create_container.assert_called_with(command='env',
                                                         environment={
                                                             'AIRFLOW_TMP_DIR': '/tmp/airflow',
                                                             'UNIT': 'TEST'
                                                         },
                                                         host_config=host_config,
                                                         image='ubuntu:latest',
-                                                        mem_limit=None, user=None,
+                                                        user=None,
                                                         working_dir='/container/path'
                                                         )
         client_mock.create_host_config.assert_called_with(binds=['/host/path:/container/path',
                                                                  '/mkdtemp:/tmp/airflow'],
                                                           network_mode='bridge',
                                                           shm_size=1000,
+                                                          cpu_shares=1024,
+                                                          mem_limit=None,
                                                           dns=None,
                                                           dns_search=None)
         client_mock.images.assert_called_with(name='ubuntu:latest')
@@ -236,5 +238,6 @@ def test_execute_with_docker_conn_id_use_hook(self, operator_client_mock,
             'Image was not pulled using operator client'
         )
 
+
 if __name__ == "__main__":
     unittest.main()

tests/operators/hive_operator.py tests/operators/test_hive_operator.py

@@ -7,9 +7,9 @@
 # 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
@@ -22,11 +22,10 @@
 import datetime
 import os
 import unittest
-import mock
 import nose
-import six
 
-from airflow import DAG, configuration, operators
+from airflow import DAG, configuration
+import airflow.operators.hive_operator
 configuration.load_test_config()
 
 
@@ -61,7 +60,7 @@ def setUp(self):
 class HiveOperatorConfigTest(HiveEnvironmentTest):
 
     def test_hive_airflow_default_config_queue(self):
-        t = operators.hive_operator.HiveOperator(
+        t = airflow.operators.hive_operator.HiveOperator(
             task_id='test_default_config_queue',
             hql=self.hql,
             mapred_queue_priority='HIGH',
@@ -77,7 +76,7 @@ def test_hive_airflow_default_config_queue(self):
 
     def test_hive_airflow_default_config_queue_override(self):
         specific_mapred_queue = 'default'
-        t = operators.hive_operator.HiveOperator(
+        t = airflow.operators.hive_operator.HiveOperator(
             task_id='test_default_config_queue',
             hql=self.hql,
             mapred_queue=specific_mapred_queue,
@@ -92,15 +91,15 @@ class HiveOperatorTest(HiveEnvironmentTest):
 
     def test_hiveconf_jinja_translate(self):
         hql = "SELECT ${num_col} FROM ${hiveconf:table};"
-        t = operators.hive_operator.HiveOperator(
+        t = airflow.operators.hive_operator.HiveOperator(
             hiveconf_jinja_translate=True,
             task_id='dry_run_basic_hql', hql=hql, dag=self.dag)
         t.prepare_template()
         self.assertEqual(t.hql, "SELECT {{ num_col }} FROM {{ table }};")
 
     def test_hiveconf(self):
         hql = "SELECT * FROM ${hiveconf:table} PARTITION (${hiveconf:day});"
-        t = operators.hive_operator.HiveOperator(
+        t = airflow.operators.hive_operator.HiveOperator(
             hiveconfs={'table': 'static_babynames', 'day': '{{ ds }}'},
             task_id='dry_run_basic_hql', hql=hql, dag=self.dag)
         t.prepare_template()
@@ -117,13 +116,13 @@ def test_hiveconf(self):
     class HivePrestoTest(HiveEnvironmentTest):
 
         def test_hive(self):
-            t = operators.hive_operator.HiveOperator(
+            t = airflow.operators.hive_operator.HiveOperator(
                 task_id='basic_hql', hql=self.hql, dag=self.dag)
             t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                   ignore_ti_state=True)
 
         def test_hive_queues(self):
-            t = operators.hive_operator.HiveOperator(
+            t = airflow.operators.hive_operator.HiveOperator(
                 task_id='test_hive_queues', hql=self.hql,
                 mapred_queue='default', mapred_queue_priority='HIGH',
                 mapred_job_name='airflow.test_hive_queues',
@@ -132,12 +131,12 @@ def test_hive_queues(self):
                   ignore_ti_state=True)
 
         def test_hive_dryrun(self):
-            t = operators.hive_operator.HiveOperator(
+            t = airflow.operators.hive_operator.HiveOperator(
                 task_id='dry_run_basic_hql', hql=self.hql, dag=self.dag)
             t.dry_run()
 
         def test_beeline(self):
-            t = operators.hive_operator.HiveOperator(
+            t = airflow.operators.hive_operator.HiveOperator(
                 task_id='beeline_hql', hive_cli_conn_id='beeline_default',
                 hql=self.hql, dag=self.dag)
             t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
@@ -147,13 +146,13 @@ def test_presto(self):
             sql = """
             SELECT count(1) FROM airflow.static_babynames_partitioned;
             """
-            t = operators.presto_check_operator.PrestoCheckOperator(
+            t = airflow.operators.presto_check_operator.PrestoCheckOperator(
                 task_id='presto_check', sql=sql, dag=self.dag)
             t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                   ignore_ti_state=True)
 
         def test_presto_to_mysql(self):
-            t = operators.presto_to_mysql.PrestoToMySqlTransfer(
+            t = airflow.operators.presto_to_mysql.PrestoToMySqlTransfer(
                 task_id='presto_to_mysql_check',
                 sql="""
                 SELECT name, count(*) as ccount
@@ -167,15 +166,15 @@ def test_presto_to_mysql(self):
                   ignore_ti_state=True)
 
         def test_hdfs_sensor(self):
-            t = operators.sensors.HdfsSensor(
+            t = airflow.operators.sensors.HdfsSensor(
                 task_id='hdfs_sensor_check',
                 filepath='hdfs://user/hive/warehouse/airflow.db/static_babynames',
                 dag=self.dag)
             t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                   ignore_ti_state=True)
 
         def test_webhdfs_sensor(self):
-            t = operators.sensors.WebHdfsSensor(
+            t = airflow.operators.sensors.WebHdfsSensor(
                 task_id='webhdfs_sensor_check',
                 filepath='hdfs://user/hive/warehouse/airflow.db/static_babynames',
                 timeout=120,
@@ -184,7 +183,7 @@ def test_webhdfs_sensor(self):
                   ignore_ti_state=True)
 
         def test_sql_sensor(self):
-            t = operators.sensors.SqlSensor(
+            t = airflow.operators.sensors.SqlSensor(
                 task_id='hdfs_sensor_check',
                 conn_id='presto_default',
                 sql="SELECT 'x' FROM airflow.static_babynames LIMIT 1;",
@@ -193,7 +192,7 @@ def test_sql_sensor(self):
                   ignore_ti_state=True)
 
         def test_hive_stats(self):
-            t = operators.hive_stats_operator.HiveStatsCollectionOperator(
+            t = airflow.operators.hive_stats_operator.HiveStatsCollectionOperator(
                 task_id='hive_stats_check',
                 table="airflow.static_babynames_partitioned",
                 partition={'ds': DEFAULT_DATE_DS},
@@ -202,7 +201,7 @@ def test_hive_stats(self):
                   ignore_ti_state=True)
 
         def test_named_hive_partition_sensor(self):
-            t = operators.sensors.NamedHivePartitionSensor(
+            t = airflow.operators.sensors.NamedHivePartitionSensor(
                 task_id='hive_partition_check',
                 partition_names=[
                     "airflow.static_babynames_partitioned/ds={{ds}}"
@@ -212,7 +211,7 @@ def test_named_hive_partition_sensor(self):
                   ignore_ti_state=True)
 
         def test_named_hive_partition_sensor_succeeds_on_multiple_partitions(self):
-            t = operators.sensors.NamedHivePartitionSensor(
+            t = airflow.operators.sensors.NamedHivePartitionSensor(
                 task_id='hive_partition_check',
                 partition_names=[
                     "airflow.static_babynames_partitioned/ds={{ds}}",
@@ -223,15 +222,15 @@ def test_named_hive_partition_sensor_succeeds_on_multiple_partitions(self):
                   ignore_ti_state=True)
 
         def test_named_hive_partition_sensor_parses_partitions_with_periods(self):
-            t = operators.sensors.NamedHivePartitionSensor.parse_partition_name(
+            t = airflow.operators.sensors.NamedHivePartitionSensor.parse_partition_name(
                 partition="schema.table/part1=this.can.be.an.issue/part2=ok")
             self.assertEqual(t[0], "schema")
             self.assertEqual(t[1], "table")
             self.assertEqual(t[2], "part1=this.can.be.an.issue/part2=this_should_be_ok")
 
         @nose.tools.raises(airflow.exceptions.AirflowSensorTimeout)
         def test_named_hive_partition_sensor_times_out_on_nonexistent_partition(self):
-            t = operators.sensors.NamedHivePartitionSensor(
+            t = airflow.operators.sensors.NamedHivePartitionSensor(
                 task_id='hive_partition_check',
                 partition_names=[
                     "airflow.static_babynames_partitioned/ds={{ds}}",
@@ -244,15 +243,15 @@ def test_named_hive_partition_sensor_times_out_on_nonexistent_partition(self):
                   ignore_ti_state=True)
 
         def test_hive_partition_sensor(self):
-            t = operators.sensors.HivePartitionSensor(
+            t = airflow.operators.sensors.HivePartitionSensor(
                 task_id='hive_partition_check',
                 table='airflow.static_babynames_partitioned',
                 dag=self.dag)
             t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                   ignore_ti_state=True)
 
         def test_hive_metastore_sql_sensor(self):
-            t = operators.sensors.MetastorePartitionSensor(
+            t = airflow.operators.sensors.MetastorePartitionSensor(
                 task_id='hive_partition_check',
                 table='airflow.static_babynames_partitioned',
                 partition_name='ds={}'.format(DEFAULT_DATE_DS),
@@ -261,7 +260,7 @@ def test_hive_metastore_sql_sensor(self):
                   ignore_ti_state=True)
 
         def test_hive2samba(self):
-            t = operators.hive_to_samba_operator.Hive2SambaOperator(
+            t = airflow.operators.hive_to_samba_operator.Hive2SambaOperator(
                 task_id='hive2samba_check',
                 samba_conn_id='tableau_samba',
                 hql="SELECT * FROM airflow.static_babynames LIMIT 10000",
@@ -271,7 +270,7 @@ def test_hive2samba(self):
                   ignore_ti_state=True)
 
         def test_hive_to_mysql(self):
-            t = operators.hive_to_mysql.HiveToMySqlTransfer(
+            t = airflow.operators.hive_to_mysql.HiveToMySqlTransfer(
                 mysql_conn_id='airflow_db',
                 task_id='hive_to_mysql_check',
                 create=True,