From 95eac511ebfd96ca62507dcde2826ad397b82555 Mon Sep 17 00:00:00 2001
From: Marcus Martins <marcus.martins@docker.com>
Date: Mon, 12 Jan 2015 13:24:37 -0800
Subject: [PATCH] Provide helper TestCases to be used with Defender

Provide TransactionTestCase and TestCase that clear the
defender cache between runs.
---
 defender/test.py  | 22 ++++++++++++++++++++++
 defender/tests.py | 42 ++++++++++++++++++++++++++++++++++--------
 2 files changed, 56 insertions(+), 8 deletions(-)
 create mode 100644 defender/test.py

diff --git a/defender/test.py b/defender/test.py
new file mode 100644
index 0000000..c72689d
--- /dev/null
+++ b/defender/test.py
@@ -0,0 +1,22 @@
+from django.test.testcases import TestCase, TransactionTestCase
+
+from .connection import get_redis_connection
+
+
+class DefenderTestCaseMixin(object):
+    """Mixin used to provide a common tearDown method"""
+
+    def tearDown(self):
+        """cleanup django-defender cache after each test"""
+        super(DefenderTestCaseMixin, self).tearDown()
+        get_redis_connection().flushdb()
+
+
+class DefenderTransactionTestCase(DefenderTestCaseMixin, TransactionTestCase):
+    """Helper TransactionTestCase that cleans the cache after each test"""
+    pass
+
+
+class DefenderTestCase(DefenderTestCaseMixin, TestCase):
+    """Helper TestCase that cleans the cache after each test"""
+    pass
diff --git a/defender/tests.py b/defender/tests.py
index 1aeefe8..18fe1d8 100644
--- a/defender/tests.py
+++ b/defender/tests.py
@@ -4,19 +4,19 @@ import time
 
 from mock import patch
 
-from django.test import TestCase
-from django.test.client import RequestFactory
 from django.contrib.auth.models import User
 from django.contrib.auth.models import AnonymousUser
 from django.contrib.sessions.backends.db import SessionStore
 from django.core.urlresolvers import NoReverseMatch
 from django.core.urlresolvers import reverse
 from django.http import HttpRequest
+from django.test.client import RequestFactory
 
-from .connection import parse_redis_url, get_redis_connection
 from . import utils
 from . import config
+from .connection import parse_redis_url, get_redis_connection
 from .models import AccessAttempt
+from .test import DefenderTestCase, DefenderTransactionTestCase
 
 
 # Django >= 1.7 compatibility
@@ -31,7 +31,7 @@ except NoReverseMatch:
 VALID_USERNAME = VALID_PASSWORD = 'valid'
 
 
-class AccessAttemptTest(TestCase):
+class AccessAttemptTest(DefenderTestCase):
     """ Test case using custom settings for testing
     """
     LOCKED_MESSAGE = 'Account locked: too many login attempts.'
@@ -74,10 +74,6 @@ class AccessAttemptTest(TestCase):
             password=VALID_PASSWORD,
         )
 
-    def tearDown(self):
-        """ clean up the db """
-        get_redis_connection().flushdb()
-
     def test_login_get(self):
         """ visit the login page """
         response = self.client.get(ADMIN_LOGIN_URL)
@@ -556,3 +552,33 @@ class AccessAttemptTest(TestCase):
         self.assertEquals(AccessAttempt.objects.count(),
                           config.FAILURE_LIMIT+1)
         self.assertIsNotNone(str(AccessAttempt.objects.all()[0]))
+
+
+class DefenderTestCaseTest(DefenderTestCase):
+    """Make sure that we're cleaning the cache between tests"""
+    key = 'test_key'
+
+    def test_first_incr(self):
+        utils.redis_server.incr(self.key)
+        result = int(utils.redis_server.get(self.key))
+        self.assertEqual(result, 1)
+
+    def test_second_incr(self):
+        utils.redis_server.incr(self.key)
+        result = int(utils.redis_server.get(self.key))
+        self.assertEqual(result, 1)
+
+
+class DefenderTransactionTestCaseTest(DefenderTransactionTestCase):
+    """Make sure that we're cleaning the cache between tests"""
+    key = 'test_key'
+
+    def test_first_incr(self):
+        utils.redis_server.incr(self.key)
+        result = int(utils.redis_server.get(self.key))
+        self.assertEqual(result, 1)
+
+    def test_second_incr(self):
+        utils.redis_server.incr(self.key)
+        result = int(utils.redis_server.get(self.key))
+        self.assertEqual(result, 1)