potiuk commented on code in PR #40511: URL: https://github.com/apache/airflow/pull/40511#discussion_r1665401566
########## tests/providers/microsoft/azure/transfers/test_s3_to_wasb.py: ########## @@ -0,0 +1,333 @@ +# +# 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. + +from __future__ import annotations + +from io import RawIOBase +from unittest import mock + +import pytest +from moto import mock_aws + +from airflow.providers.microsoft.azure.transfers.s3_to_wasb import S3ToAzureBlobStorageOperator + +TASK_ID = "test-s3-to-azure-blob-operator" +AWS_CONN_ID = "test-conn-id" +S3_BUCKET = "test-bucket" +CONTAINER_NAME = "test-container" +PREFIX = "TEST" +TEMPFILE_NAME = "test-tempfile" +MOCK_FILES = ["TEST1.csv", "TEST2.csv", "TEST3.csv"] + +# Here are some of the tests that need to be run (for the get_files_to_move() function) +# 1. Prefix with no existing files, without replace [DONE] +# 2. Prefix with existing files, without replace [DONE] +# 3. Prefix with existing files, with replace [DONE] +# 4. Two keys without existing files, without replace [DONE] +# 5. Two keys with existing files, without replace [DONE] +# 6. Two keys with existing files, with replace [DONE] +# 7. S3 key with Azure prefix, without existing files, without replace [DONE] +# 8. S3 key with Azure prefix, with existing files, without replace [DONE] +# 9. S3 key with Azure prefix, with existing files, with replace [SKIPPED] + +# Other tests that need to be run +# - Test __init__ [DONE] +# - Test execute() [DONE] + + +@mock_aws +class TestS3ToAzureBlobStorageOperator: + def test__init__(self): + # Create a mock operator with a single set of parameters that are used to test the __init__() + # constructor. Not every parameter needs to be provided, as this code will also be used to test the + # default parameters that are configured + operator = S3ToAzureBlobStorageOperator( + task_id=TASK_ID, + aws_conn_id="test-conn-id", + s3_bucket=S3_BUCKET, + s3_prefix=PREFIX, + container_name=CONTAINER_NAME, + blob_prefix=PREFIX, + replace=True, + ) + + # ... is None is used to validate if a value is None, while not ... is used to evaluate if a value + # is False + assert operator.task_id == TASK_ID + assert operator.aws_conn_id == AWS_CONN_ID + assert operator.wasb_conn_id == "wasb_default" + assert operator.s3_bucket == S3_BUCKET + assert operator.container_name == CONTAINER_NAME + assert operator.s3_prefix == PREFIX + assert operator.s3_key is None + assert operator.blob_prefix == PREFIX + assert operator.blob_name is None + assert not operator.create_container # Should be false (match default value in constructor) + assert operator.replace + assert not operator.s3_verify # Should be false (match default value in constructor) + assert operator.s3_extra_args == {} + assert operator.wasb_extra_args == {} + + @mock.patch("airflow.providers.microsoft.azure.transfers.s3_to_wasb.S3Hook") + @mock.patch("airflow.providers.microsoft.azure.transfers.s3_to_wasb.WasbHook") + @mock.patch("tempfile.NamedTemporaryFile") + def test__execute__prefix_without_replace_empty_destination( + self, tempfile_mock, wasb_mock_hook, s3_mock_hook + ): + # Set the list files that the S3Hook should return, along with an empty list of files in the Azure + # Blob storage container. This scenario was picked for testing, as it's most likely the most common + # setting the operator will be used in + s3_mock_hook.return_value.list_keys.return_value = MOCK_FILES + wasb_mock_hook.return_value.get_blobs_list_recursive.return_value = [] + + s3_mock_hook.return_value.download_file.return_value = RawIOBase(b"test file contents") + tempfile_mock.return_value.__enter__.return_value.name = TEMPFILE_NAME + + operator = S3ToAzureBlobStorageOperator( + task_id=TASK_ID, + s3_bucket=S3_BUCKET, + s3_prefix=PREFIX, + container_name=CONTAINER_NAME, + blob_prefix=PREFIX, + ) + # Placing an empty "context" object here (using None) + uploaded_files = operator.execute(None) + assert sorted(uploaded_files) == sorted(MOCK_FILES) + + # Using the default connection ID, along with the default value of verify (for the S3 hook) + s3_mock_hook.assert_called_once_with(aws_conn_id="aws_default", verify=False) + wasb_mock_hook.assert_called_once_with(wasb_conn_id="wasb_default") + + # There are a number of very similar tests that use the same mocking, and for the most part, the same Review Comment: Ah right. My bad. Can you please inline the variables instead of having constants ? I was mislead by that - when you use parametrise it's best to just have the array of parameters inside the decorator, then you don't have to move up to the constant declaration to see what are the paremeters for testing. Also Doc build needs to be fixed. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
