From b18f70205c54cb31c9947368d8bb4627193b71c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lubo=C5=A1=20M=C3=A1tl?= <matllubos@gmail.com>
Date: Tue, 21 Dec 2021 13:50:26 +0100
Subject: [PATCH] Faster and bullet proof test runner

---
 pydjamodb/connection.py          | 33 ++++++++++++++++++++++++++++++++
 pydjamodb/test_runner.py         | 23 +++++++++++++++++-----
 pydjamodb/tests.py               | 11 -----------
 tests/requirements.txt           |  2 +-
 tests/test_app/tests/__init__.py |  5 ++---
 5 files changed, 54 insertions(+), 20 deletions(-)
 delete mode 100644 pydjamodb/tests.py

diff --git a/pydjamodb/connection.py b/pydjamodb/connection.py
index 4485a78..f97262a 100644
--- a/pydjamodb/connection.py
+++ b/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
diff --git a/pydjamodb/test_runner.py b/pydjamodb/test_runner.py
index 676635a..90dd436 100644
--- a/pydjamodb/test_runner.py
+++ b/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)
diff --git a/pydjamodb/tests.py b/pydjamodb/tests.py
deleted file mode 100644
index 93c18dc..0000000
--- a/pydjamodb/tests.py
+++ /dev/null
@@ -1,11 +0,0 @@
-from .models import dynamodb_model_classes
-
-
-class DynamoDBTestMixin:
-
-    def tearDown(self):
-        super().tearDown()
-        for model_class in dynamodb_model_classes:
-            with model_class.batch_write() as batch:
-                for item in model_class.scan():
-                    batch.delete(item)
diff --git a/tests/requirements.txt b/tests/requirements.txt
index ad89cbc..2a3869a 100644
--- a/tests/requirements.txt
+++ b/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
diff --git a/tests/test_app/tests/__init__.py b/tests/test_app/tests/__init__.py
index aaf7302..3f31f9c 100644
--- a/tests/test_app/tests/__init__.py
+++ b/tests/test_app/tests/__init__.py
@@ -2,9 +2,9 @@
 
 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
@@ -12,10 +12,9 @@
 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(