test: Added cleanup of old sink storage buckets

Author: gkevinzheng

samples/snippets/export_test.py

@@ -19,7 +19,7 @@
 import time
 
 import backoff
-from google.cloud import logging
+from google.cloud import logging, storage
 import pytest
 
 import export
@@ -46,8 +46,8 @@ def _create_sink_name():
 
 
 @backoff.on_exception(backoff.expo, Exception, max_time=60, raise_on_giveup=False)
-def _delete_sink(sink):
-    sink.delete()
+def _delete_object(obj):
+    obj.delete()
 
 
 # Runs once for entire test suite
@@ -62,7 +62,20 @@ def cleanup_old_sinks():
         if match:
             sink_timestamp = int(match.group(1))
             if TIMESTAMP - sink_timestamp > CLEANUP_THRESHOLD:
-                _delete_sink(sink)
+                _delete_object(sink)
+    
+    storage_client = storage.Client()
+
+    # See _sink_storage_setup in usage_guide.py for details about how
+    # sinks are named.
+    test_bucket_name_regex = r"^sink\-storage\-(\d+)$"
+    for bucket in storage_client.list_buckets():
+        match = re.match(test_bucket_name_regex, bucket.name)
+        if match:
+            # Bucket timestamp is int(time.time() * 1000)
+            bucket_timestamp = int(match.group(1))
+            if TIMESTAMP - bucket_timestamp//1000 > CLEANUP_THRESHOLD:
+                _delete_object(bucket)
 
 
 @pytest.fixture
@@ -79,7 +92,7 @@ def example_sink(cleanup_old_sinks):
 
     yield sink
 
-    _delete_sink(sink)
+    _delete_object(sink)
 
 
 def test_list(example_sink, capsys):
@@ -99,7 +112,7 @@ def test_create(capsys):
         export.create_sink(sink_name, BUCKET, TEST_SINK_FILTER)
     # Clean-up the temporary sink.
     finally:
-        _delete_sink(logging.Client().sink(sink_name))
+        _delete_object(logging.Client().sink(sink_name))
 
     out, _ = capsys.readouterr()
     assert sink_name in out