Convert Athena Sample DAG to System Test (#24058)

docs/apache-airflow-providers-amazon/example-dags.rst

@@ -0,0 +1,23 @@
+ .. 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.
+
+Example DAGs
+============
+
+You can learn how to use Amazon AWS integrations by analyzing the source code of the example DAGs:
+
+* `Amazon AWS <https://github.com/apache/airflow/tree/main/airflow/providers/amazon/aws/example_dags>`__

docs/apache-airflow-providers-amazon/index.rst

@@ -40,7 +40,7 @@ Content
     :maxdepth: 1
     :caption: Resources
 
-    Example DAGs <https://github.com/apache/airflow/tree/main/airflow/providers/amazon/aws/example_dags>
+    Example DAGs <example-dags>
     PyPI Repository <https://pypi.org/project/apache-airflow-providers-amazon/>
     Installing from sources <installing-providers-from-sources>
 

docs/apache-airflow-providers-amazon/operators/athena.rst

@@ -45,7 +45,7 @@ In the following example, we query an existing Athena table and send the results
 an existing Amazon S3 bucket.  For more examples of how to use this operator, please
 see the `Sample DAG <https://github.com/apache/airflow/blob/main/airflow/providers/amazon/aws/example_dags/example_athena.py>`__.
 
-.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_athena.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_athena.py
     :language: python
     :start-after: [START howto_operator_athena]
     :dedent: 4
@@ -62,7 +62,7 @@ Wait on Amazon Athena query results
 Use the :class:`~airflow.providers.amazon.aws.sensors.athena.AthenaSensor`
 to wait for the results of a query in Amazon Athena.
 
-.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_athena.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_athena.py
     :language: python
     :start-after: [START howto_sensor_athena]
     :dedent: 4

tests/providers/amazon/aws/system/__init__.py

@@ -0,0 +1,16 @@
+# 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.

tests/providers/amazon/aws/system/utils/__init__.py

@@ -0,0 +1,16 @@
+# 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.

tests/providers/amazon/aws/system/utils/test_helpers.py

@@ -0,0 +1,144 @@
+# 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.
+
+"""
+This module contains the unit tests for the helper methods included in the Amazon System Tests found at
+tests/system/providers/amazon/aws/utils/__init__.py
+"""
+import io
+import os
+import sys
+from unittest.mock import ANY, patch
+
+import pytest
+
+from tests.system.providers.amazon.aws import utils
+from tests.system.providers.amazon.aws.utils import (
+    DEFAULT_ENV_ID_LEN,
+    DEFAULT_ENV_ID_PREFIX,
+    ENV_ID_ENVIRON_KEY,
+    INVALID_ENV_ID_MSG,
+    LOWERCASE_ENV_ID_MSG,
+    NO_VALUE_MSG,
+    _validate_env_id,
+    set_env_id,
+)
+
+try:
+    from moto import mock_ssm
+except ImportError:
+    mock_ssm = None
+
+TEST_NAME: str = 'example_test'
+ANY_STR: str = 'any'
+
+ENV_VALUE: str = 'foo'
+SSM_VALUE: str = 'bar'
+DEFAULT_VALUE: str = 'baz'
+
+
+@pytest.fixture(autouse=True)
+def provide_test_name():
+    with patch.object(utils, '_get_test_name', return_value=TEST_NAME) as name:
+        yield name
+
+
+@pytest.mark.skipif(mock_ssm is None, reason='mock_ssm package not present')
+@mock_ssm
+class TestAmazonSystemTestHelpers:
+    FETCH_VARIABLE_TEST_CASES = [
+        # Format is:
+        # (Environment Variable value, Fetched SSM value, Provided Default value, Expected Result)
+        (ENV_VALUE, SSM_VALUE, DEFAULT_VALUE, ENV_VALUE),
+        (ENV_VALUE, SSM_VALUE, None, ENV_VALUE),
+        (ENV_VALUE, None, DEFAULT_VALUE, ENV_VALUE),
+        (ENV_VALUE, None, None, ENV_VALUE),
+        (None, SSM_VALUE, DEFAULT_VALUE, SSM_VALUE),
+        (None, SSM_VALUE, None, SSM_VALUE),
+        (None, None, DEFAULT_VALUE, DEFAULT_VALUE),
+        # For the (None, None, None ) test case, see: test_fetch_variable_no_value_found_raises_exception
+    ]
+
+    @pytest.mark.parametrize(
+        'env_value, ssm_value, default_value, expected_result', FETCH_VARIABLE_TEST_CASES
+    )
+    @patch.object(os, 'getenv')
+    def test_fetch_variable_success(
+        self, mock_getenv, env_value, ssm_value, default_value, expected_result
+    ) -> None:
+        mock_getenv.return_value = env_value if env_value else ssm_value
+
+        result = utils.fetch_variable(ANY, default_value) if default_value else utils.fetch_variable(ANY_STR)
+
+        assert result == expected_result
+
+    def test_fetch_variable_no_value_found_raises_exception(self):
+        # This would be the (None, None, None) test case from above.
+        with pytest.raises(ValueError) as raised_exception:
+            utils.fetch_variable(ANY_STR)
+
+            assert NO_VALUE_MSG.format(key=ANY_STR) in str(raised_exception.value)
+
+    ENV_ID_TEST_CASES = [
+        # Happy Cases
+        ('ABCD', True),
+        ('AbCd', True),
+        ('abcd', True),
+        ('ab12', True),
+        # Failure Cases
+        # Must be alphanumeric
+        ('not_alphanumeric', False),
+        # Can not be empty
+        ('', False),
+        # Must start with a letter
+        ('1234', False),
+        ('12ab', False),
+        ('12AB', False),
+        ('12Ab', False),
+    ]
+
+    @pytest.mark.parametrize('env_id, is_valid', ENV_ID_TEST_CASES)
+    def test_validate_env_id_success(self, env_id, is_valid):
+        if is_valid:
+            captured_output = io.StringIO()
+            sys.stdout = captured_output
+
+            result = _validate_env_id(env_id)
+            sys.stdout = sys.__stdout__
+
+            assert result == env_id.lower()
+            assert result.isalnum()
+            if not result == env_id:
+                assert LOWERCASE_ENV_ID_MSG in captured_output.getvalue()
+        else:
+            with pytest.raises(ValueError) as raised_exception:
+                _validate_env_id(env_id)
+
+                assert INVALID_ENV_ID_MSG in str(raised_exception.value)
+
+    def test_set_env_id_generates_if_required(self):
+        # No environment variable nor SSM value has been found
+        result = set_env_id()
+
+        assert len(result) == DEFAULT_ENV_ID_LEN + len(DEFAULT_ENV_ID_PREFIX)
+        assert result.isalnum()
+        assert result.islower()
+
+    def test_set_env_id_exports_environment_variable(self):
+        env_id = set_env_id()
+
+        assert os.environ[ENV_ID_ENVIRON_KEY] == env_id

tests/system/providers/amazon/README.md

@@ -0,0 +1,77 @@
+<!--
+ 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.
+-->
+
+# Amazon provider system tests
+
+## Tests structure
+
+All AWS-related system tests are located inside `tests/system/providers/amazon/aws/`.
+In this directory you will find test files in the form of Example DAGs, one DAG per file.
+Each test should be self-contained but in the case where additional resources are required,
+they can be found in the `resources` directory on the same level as tests or noted in the
+test's docstring.  Each test file should start with prefix `example_*`.
+
+Example directory structure:
+
+```
+tests/system/providers/amazon/aws
+├── resources
+│   ├── example_athena_data.csv
+│   └── example_sagemaker_constants.py
+├── example_athena.py
+├── example_batch.py
+.
+├── example_step_functions.py
+└── *
+```
+
+## Initial configuration
+
+Each test requires some environment variables. Check how to set them up on your
+operating system, but on UNIX-based operating systems `export NAME_OF_ENV_VAR=value`
+should work.  To confirm that it is set up correctly, run `echo $NAME_OF_ENV_VAR`
+which will display its value.
+
+When manually running tests using pytest, you can define them inline with the command.
+For example:
+
+```commandline
+NAME_OF_ENV_VAR=value pytest --system amazon tests/system/providers/amazon/aws/example_test.py
+```
+
+### Required environment variables
+
+- `SYSTEM_TESTS_ENV_ID` - AWS System Tests will generate and export this value if one does not exist.
+
+  An environment ID is a unique value across different executions of system tests.  This
+  is needed because the CI environment may run the tests on various versions of Airflow
+  in parallel.  If this is the case, the value of this variable ensures that resources
+  that are created during the tests will not interfere with each other.
+
+  The value is used as part of the name for resources which have different requirements.
+  For example: an S3 bucket name can not use underscores, but an Athena table name can not
+  use hyphens.  In order to minimize conflicts, this variable should be a randomized value
+  using only lowercase letters A-Z and digits 0-9, and start with a letter.
+
+## Settings for specific tests
+
+Amazon system test files are designed to be as self-contained as possible.  They will contain
+any sample data and configuration values which are required, and they will create and tear
+down any required infrastructure.  Some tests will require an IAM Role ARN, and the requirements
+for those Roles should be documented inside the docstring of the test file.