[AIRFLOW-3141] Add missing missing sensor tests. (apache#3991)

airflow/sensors/s3_key_sensor.py

@@ -73,10 +73,7 @@ def __init__(self,
                 raise AirflowException('Please provide a bucket_name')
             else:
                 bucket_name = parsed_url.netloc
-                if parsed_url.path[0] == '/':
-                    bucket_key = parsed_url.path[1:]
-                else:
-                    bucket_key = parsed_url.path
+                bucket_key = parsed_url.path.lstrip('/')
         else:
             parsed_url = urlparse(bucket_key)
             if parsed_url.scheme != '' or parsed_url.netloc != '':
@@ -97,5 +94,4 @@ def poke(self, context):
         if self.wildcard_match:
             return hook.check_for_wildcard_key(self.bucket_key,
                                                self.bucket_name)
-        else:
-            return hook.check_for_key(self.bucket_key, self.bucket_name)
+        return hook.check_for_key(self.bucket_key, self.bucket_name)

airflow/sensors/sql_sensor.py

@@ -51,8 +51,4 @@ def poke(self, context):
         records = hook.get_records(self.sql)
         if not records:
             return False
-        else:
-            if str(records[0][0]) in ('0', '',):
-                return False
-            else:
-                return True
+        return str(records[0][0]) not in ('0', '')

tests/sensors/test_s3_key_sensor.py

@@ -17,7 +17,10 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import mock
 import unittest
+from parameterized import parameterized
+
 from airflow.exceptions import AirflowException
 from airflow.sensors.s3_key_sensor import S3KeySensor
 
@@ -31,7 +34,9 @@ def test_bucket_name_None_and_bucket_key_as_relative_path(self):
         :return:
         """
         with self.assertRaises(AirflowException):
-            S3KeySensor(bucket_key="file_in_bucket")
+            S3KeySensor(
+                task_id='s3_key_sensor',
+                bucket_key="file_in_bucket")
 
     def test_bucket_name_provided_and_bucket_key_is_s3_url(self):
         """
@@ -40,5 +45,49 @@ def test_bucket_name_provided_and_bucket_key_is_s3_url(self):
         :return:
         """
         with self.assertRaises(AirflowException):
-            S3KeySensor(bucket_key="s3://test_bucket/file",
-                        bucket_name='test_bucket')
+            S3KeySensor(
+                task_id='s3_key_sensor',
+                bucket_key="s3://test_bucket/file",
+                bucket_name='test_bucket')
+
+    @parameterized.expand([
+        ['s3://bucket/key', None, 'key', 'bucket'],
+        ['key', 'bucket', 'key', 'bucket'],
+    ])
+    def test_parse_bucket_key(self, key, bucket, parsed_key, parsed_bucket):
+        s = S3KeySensor(
+            task_id='s3_key_sensor',
+            bucket_key=key,
+            bucket_name=bucket,
+        )
+        self.assertEqual(s.bucket_key, parsed_key)
+        self.assertEqual(s.bucket_name, parsed_bucket)
+
+    @mock.patch('airflow.hooks.S3_hook.S3Hook')
+    def test_poke(self, mock_hook):
+        s = S3KeySensor(
+            task_id='s3_key_sensor',
+            bucket_key='s3://test_bucket/file')
+
+        mock_check_for_key = mock_hook.return_value.check_for_key
+        mock_check_for_key.return_value = False
+        self.assertFalse(s.poke(None))
+        mock_check_for_key.assert_called_with(s.bucket_key, s.bucket_name)
+
+        mock_hook.return_value.check_for_key.return_value = True
+        self.assertTrue(s.poke(None))
+
+    @mock.patch('airflow.hooks.S3_hook.S3Hook')
+    def test_poke_wildcard(self, mock_hook):
+        s = S3KeySensor(
+            task_id='s3_key_sensor',
+            bucket_key='s3://test_bucket/file',
+            wildcard_match=True)
+
+        mock_check_for_wildcard_key = mock_hook.return_value.check_for_wildcard_key
+        mock_check_for_wildcard_key.return_value = False
+        self.assertFalse(s.poke(None))
+        mock_check_for_wildcard_key.assert_called_with(s.bucket_key, s.bucket_name)
+
+        mock_check_for_wildcard_key.return_value = True
+        self.assertTrue(s.poke(None))

tests/sensors/test_s3_prefix_sensor.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+#
+# 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 mock
+import unittest
+
+from airflow.sensors.s3_prefix_sensor import S3PrefixSensor
+
+
+class S3PrefixSensorTests(unittest.TestCase):
+
+    @mock.patch('airflow.hooks.S3_hook.S3Hook')
+    def test_poke(self, mock_hook):
+        s = S3PrefixSensor(
+            task_id='s3_prefix',
+            bucket_name='bucket',
+            prefix='prefix')
+
+        mock_hook.return_value.check_for_prefix.return_value = False
+        self.assertFalse(s.poke(None))
+        mock_hook.return_value.check_for_prefix.assert_called_with(
+            prefix='prefix',
+            delimiter='/',
+            bucket_name='bucket')
+
+        mock_hook.return_value.check_for_prefix.return_value = True
+        self.assertTrue(s.poke(None))

tests/sensors/test_sql_sensor.py

@@ -16,6 +16,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import mock
 import unittest
 
 from airflow import DAG
@@ -30,6 +31,7 @@
 
 
 class SqlSensorTests(unittest.TestCase):
+
     def setUp(self):
         configuration.load_test_config()
         args = {
@@ -55,3 +57,21 @@ def test_sql_sensor_postgres(self):
             dag=self.dag
         )
         t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
+
+    @mock.patch('airflow.sensors.sql_sensor.BaseHook')
+    def test_sql_sensor_postgres_poke(self, mock_hook):
+        t = SqlSensor(
+            task_id='sql_sensor_check',
+            conn_id='postgres_default',
+            sql="SELECT 1",
+        )
+
+        mock_get_records = (
+            mock_hook.get_connection.return_value
+            .get_hook.return_value.get_records)
+
+        mock_get_records.return_value = []
+        self.assertFalse(t.poke(None))
+
+        mock_get_records.return_value = [['1']]
+        self.assertTrue(t.poke(None))