Add TemporaryRegistry helper

Author: cdebost

src/aws_schema_registry/client.py

@@ -1,6 +1,9 @@
+from datetime import datetime
 import logging
 import time
-from typing import Mapping
+import random
+import string
+from typing import ContextManager, Mapping
 from uuid import UUID
 
 from aws_schema_registry.schema import (
@@ -310,3 +313,57 @@ def create_schema(
                     f'Create schema {name} failed'
                 ) from e
         return version_id
+
+
+class TemporaryRegistry(ContextManager):
+    """A real schema registry for use in tests and experiments.
+
+    This class implements the ContextManager protocol, creating the registry
+    on enter and destroying it on exit.
+
+    Usage:
+
+```python
+    with TemporaryRegistry(glue_client, 'MyRegistry') as r:
+        # registry is created on enter
+        print(r.name)  # the "real" (suffixed) registry name
+        # registry is destroyed on exit
+```
+
+    Arguments:
+        glue_client: glue client created by `botocore`/`boto3`.
+        name: human-readable name for the created registry. The name will be
+            suffixed by a random identifier to reduce the freqency of
+            collisions.
+        description: description for the created registry.
+        autoremove: whether to destroy the created registry. Defaults to True.
+    """
+
+    DEFAULT_DESCRIPTION = 'Temporary registry created with the aws-glue-schema-registry Python library.'  # NOQA
+
+    def __init__(self, glue_client,
+                 name: str = 'temporary-registry',
+                 description: str = DEFAULT_DESCRIPTION,
+                 autoremove: bool = True):
+        self.glue_client = glue_client
+        date = datetime.utcnow().strftime('%y-%m-%d-%H-%M')
+        r = ''.join(random.choices(string.digits + string.ascii_letters,
+                                   k=16))
+        self.name = f'{name}-{date}-{r}'
+        self.description = description
+        self.autoremove = autoremove
+
+    def __enter__(self):
+        LOG.info('creating registry %s...' % self.name)
+        self.glue_client.create_registry(
+            RegistryName=self.name,
+            Description=self.description
+        )
+        return self
+
+    def __exit__(self, *args):
+        if self.autoremove:
+            LOG.info('deleting registry %s...' % self.name)
+            self.glue_client.delete_registry(
+                RegistryId={'RegistryName': self.name}
+            )

tests/integration/conftest.py

@@ -1,10 +1,11 @@
-from datetime import datetime
 import logging
 import os
 
 import boto3
 import pytest
 
+from aws_schema_registry.client import TemporaryRegistry
+
 LOG = logging.getLogger(__name__)
 
 AWS_ACCESS_KEY_ID = os.getenv('AWS_ACCESS_KEY_ID')
@@ -13,8 +14,6 @@
 AWS_REGION = os.getenv('AWS_REGION')
 AWS_PROFILE = os.getenv('AWS_PROFILE')
 
-REGISTRY_PREFIX = 'schema-registry-tests'
-DATE = datetime.utcnow().strftime('%y-%m-%d-%H-%M-%s')
 CLEANUP_REGISTRY = os.getenv('CLEANUP_REGISTRY') is None
 
 
@@ -37,15 +36,12 @@ def glue_client(boto_session):
 @pytest.fixture(scope='session')
 def registry(glue_client):
     """The AWS Glue registry to use for testing."""
-    name = f'{REGISTRY_PREFIX}-{DATE}'
-    LOG.info('creating registry %s...' % name)
-    glue_client.create_registry(
-        RegistryName=name,
-        Description='Registry used for the schema registry python integration'
+    with TemporaryRegistry(
+        glue_client,
+        name='schema-registry-tests',
+        description='Registry used for the schema registry python integration'
         ' tests. This registry does not hold any valuable data and is safe to'
-        ' delete as long as it is not currently in use by a test'
-    )
-    yield name
-    if CLEANUP_REGISTRY:
-        LOG.info('deleting registry %s...' % name)
-        glue_client.delete_registry(RegistryId={'RegistryName': name})
+        ' delete as long as it is not currently in use by a test',
+        autoremove=CLEANUP_REGISTRY
+    ) as registry:
+        yield registry.name