Merge pull request #8 from druids/FasterTestRunner

Faster and bullet proof test runner

pydjamodb/connection.py

@@ -139,3 +139,36 @@ def set_point_in_time_recovery(self, enabled=True):
                 if i == self.connection._max_retry_attempts_exception:
                     raise
                 time.sleep(10)
+
+
+class TestTableConnection:
+
+    def __init__(self, wrapped_connection, prefix=None):
+        self._wrapped_connection = wrapped_connection
+        self._is_test_clean_required = False
+        if prefix:
+            self._wrapped_connection.table_name = 'test_{}_{}'.format(prefix, self._wrapped_connection.table_name)
+        else:
+            self._wrapped_connection.table_name = 'test_{}'.format(self._wrapped_connection.table_name)
+
+    def __getattr__(self, attr):
+        return getattr(self._wrapped_connection, attr)
+
+    def update_item(self, *args, **kwargs):
+        self._is_test_clean_required = True
+        return self._wrapped_connection.update_item(*args, **kwargs)
+
+    def put_item(self, *args, **kwargs):
+        self._is_test_clean_required = True
+        return self._wrapped_connection.put_item(*args, **kwargs)
+
+    def batch_write_item(self, *args, **kwargs):
+        self._is_test_clean_required = True
+        return self._wrapped_connection.batch_write_item(*args, **kwargs)
+
+    def post_test_clean(self, model_class):
+        if self._is_test_clean_required:
+            with model_class.batch_write() as batch:
+                for item in model_class.scan():
+                    batch.delete(item)
+        self._is_test_clean_required = False

pydjamodb/test_runner.py

@@ -4,17 +4,24 @@
 from django.db import connections
 from django.test.runner import ParallelTestSuite, DiscoverRunner
 
+from .connection import TestTableConnection
 from .models import dynamodb_model_classes
 
+try:
+    from germanium.signals import set_up, tear_down
+
+    def set_dynamodb_test_autoclean():
+        set_up.connect(clean_dynamodb_database)
+        tear_down.connect(clean_dynamodb_database)
+except ImportError:
+    def set_dynamodb_test_autoclean():
+        pass
+
 
 def init_pynamodb_test_prefix(prefix=None):
     for model_class in dynamodb_model_classes:
         model_class._connection = None
-        connection = model_class._get_connection()
-        if prefix:
-            connection.table_name = 'test_{}_{}'.format(prefix, connection.table_name)
-        else:
-            connection.table_name = 'test_{}'.format(connection.table_name)
+        model_class._connection = TestTableConnection(model_class._get_connection(), prefix)
 
 
 def remove_pynamodb_table(model_class):
@@ -55,6 +62,11 @@ class DynamoDBParallelTestSuite(ParallelTestSuite):
     init_worker = _init_worker
 
 
+def clean_dynamodb_database(sender, **kwargs):
+    for model_class in dynamodb_model_classes:
+        model_class._connection.post_test_clean(model_class)
+
+
 class DynamoDBTestSuiteMixin:
 
     parallel_test_suite = DynamoDBParallelTestSuite
@@ -82,6 +94,7 @@ def teardown_databases(self, old_config, **kwargs):
 
     def _setup_pynamodb_database(self, prefix=None):
         init_pynamodb_test_prefix(prefix)
+        set_dynamodb_test_autoclean()
         table_names = []
         for model_class in dynamodb_model_classes:
             table_names.append(model_class._connection.table_name)

tests/requirements.txt

@@ -1,5 +1,5 @@
 django>=2.0, <4.0
 https://github.com/druids/pynamodb/tarball/AddTagsSupport#egg=pynamodb
-django-germanium==2.2.3
+https://github.com/druids/germanium/tarball/AddedSignals
 flake8
 coveralls

tests/test_app/tests/__init__.py

@@ -2,20 +2,19 @@
 
 import string
 
-from django.test import TestCase
 from django.utils.timezone import now
 
+from germanium.test_cases.default import GermaniumTestCase
 from germanium.tools import assert_equal, assert_raises, assert_true, assert_false
 
 from uuid import uuid4
 
 from test_app.models import TestDynamoModel
 
 from pydjamodb.queryset import DynamoDBQuerySetError, MultipleObjectsReturned, ObjectDoesNotExist
-from pydjamodb.tests import DynamoDBTestMixin
 
 
-class PyDjamoDBTestCase(DynamoDBTestMixin, TestCase):
+class PyDjamoDBTestCase(GermaniumTestCase):
 
     def create_test_dynamo_model(self, **kwargs):
         default_data = dict(