From 63984b7d366b647ef45fce8db2cfa00382fbf001 Mon Sep 17 00:00:00 2001
From: GitHub Action <tomas@tomasfarias.dev>
Date: Sat, 22 Jan 2022 12:43:11 +0100
Subject: [PATCH] feat(test): Clean mocked S3 bucket after every use

---
 tests/conftest.py                  | 32 +++++++++++++++++++++++++++---
 tests/hooks/s3/test_dbt_s3_hook.py | 21 --------------------
 2 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 0e8a581..81f68a7 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -9,6 +9,7 @@
 from pytest_postgresql.janitor import DatabaseJanitor
 
 from airflow_dbt_python.hooks.dbt import DbtHook
+from airflow_dbt_python.hooks.s3 import DbtS3Hook
 
 PROFILES = """
 default:
@@ -234,11 +235,36 @@ def mocked_s3_res():
 
 
 @pytest.fixture
-def s3_bucket(mocked_s3_res):
-    """Return a mocked s3 bucket for testing."""
+def s3_hook():
+    """Provide a DbtHook."""
+    return DbtS3Hook()
+
+
+@pytest.fixture
+def s3_bucket(mocked_s3_res, s3_hook):
+    """Return a mocked s3 bucket for testing.
+
+    Bucket is cleaned after every use.
+    """
     bucket = "airflow-dbt-test-s3-bucket"
     mocked_s3_res.create_bucket(Bucket=bucket)
-    return bucket
+
+    yield bucket
+
+    keys = s3_hook.list_keys(
+        bucket,
+        "",
+    )
+    if keys is not None and len(keys) > 0:
+        s3_hook.delete_objects(
+            bucket,
+            keys,
+        )
+        keys = s3_hook.list_keys(
+            bucket,
+            "",
+        )
+    assert keys is None or len(keys) == 0
 
 
 BROKEN_SQL = """
diff --git a/tests/hooks/s3/test_dbt_s3_hook.py b/tests/hooks/s3/test_dbt_s3_hook.py
index 18c2047..ce6548a 100644
--- a/tests/hooks/s3/test_dbt_s3_hook.py
+++ b/tests/hooks/s3/test_dbt_s3_hook.py
@@ -311,29 +311,10 @@ def test_push_dbt_project_to_zip_file(s3_bucket, tmpdir, test_files):
     assert key is True
 
 
-def clean_s3_prefix(hook: DbtS3Hook, prefix: str, s3_bucket: str):
-    """Ensure we are working with an empty S3 prefix."""
-    keys = hook.list_keys(
-        s3_bucket,
-        prefix,
-    )
-    if keys is not None and len(keys) > 0:
-        hook.delete_objects(
-            s3_bucket,
-            keys,
-        )
-        keys = hook.list_keys(
-            s3_bucket,
-            prefix,
-        )
-    assert keys is None or len(keys) == 0
-
-
 def test_push_dbt_project_to_files(s3_bucket, tmpdir, test_files):
     """Test pushing a dbt project to a S3 path."""
     hook = DbtS3Hook()
     prefix = f"s3://{s3_bucket}/project/"
-    clean_s3_prefix(hook, prefix, s3_bucket)
 
     hook.push_dbt_project(f"s3://{s3_bucket}/project/", test_files[0].parent.parent)
     keys = hook.list_keys(
@@ -352,7 +333,6 @@ def test_push_dbt_project_with_no_replace(s3_bucket, tmpdir, test_files):
 
     hook = DbtS3Hook()
     prefix = f"s3://{s3_bucket}/project/"
-    clean_s3_prefix(hook, prefix, s3_bucket)
     bucket = hook.get_bucket(s3_bucket)
 
     last_modified_expected = {}
@@ -408,7 +388,6 @@ def test_push_dbt_project_with_partial_replace(s3_bucket, tmpdir, test_files):
 
     hook = DbtS3Hook()
     prefix = f"s3://{s3_bucket}/project/"
-    clean_s3_prefix(hook, prefix, s3_bucket)
     bucket = hook.get_bucket(s3_bucket)
 
     last_modified_expected = {}