From a4f766fd144b8dcbf86effc2255290b3b420228c Mon Sep 17 00:00:00 2001
From: Bertrand Bordage <bordage.bertrand@gmail.com>
Date: Mon, 20 Oct 2014 17:25:20 +0200
Subject: [PATCH] Adds tests for select_for_update.

---
 cachalot/tests/read.py  | 24 ++++++++++++++++++++----
 cachalot/tests/write.py | 25 +++++++++++++++++++++++--
 2 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/cachalot/tests/read.py b/cachalot/tests/read.py
index 5ee7870..f70dfed 100644
--- a/cachalot/tests/read.py
+++ b/cachalot/tests/read.py
@@ -8,9 +8,10 @@ except ImportError:  # For Python 2.6
     from unittest2 import skip
 
 from django.contrib.auth.models import Group, Permission, User
-from django.db import connection
+from django.db import connection, transaction
 from django.db.models import Count
-from django.test import TransactionTestCase
+from django.db.transaction import TransactionManagementError
+from django.test import TransactionTestCase, skipUnlessDBFeature
 
 from .models import Test
 
@@ -439,9 +440,24 @@ class ReadTestCase(TransactionTestCase):
     def test_using(self):
         pass
 
-    @skip(NotImplementedError)
+    @skipUnlessDBFeature('has_select_for_update')
     def test_select_for_update(self):
-        pass
+        with self.assertRaises(TransactionManagementError):
+            list(Test.objects.select_for_update())
+
+        with self.assertNumQueries(1):
+            with transaction.atomic():
+                data1 = list(Test.objects.select_for_update())
+                self.assertListEqual(data1, [self.t1, self.t2])
+                self.assertListEqual([t.name for t in data1],
+                                     ['test1', 'test2'])
+
+        with self.assertNumQueries(0):
+            with transaction.atomic():
+                data2 = list(Test.objects.select_for_update())
+                self.assertListEqual(data2, [self.t1, self.t2])
+                self.assertListEqual([t.name for t in data2],
+                                     ['test1', 'test2'])
 
     def test_extra_select(self):
         """
diff --git a/cachalot/tests/write.py b/cachalot/tests/write.py
index 0b69c9b..221eb98 100644
--- a/cachalot/tests/write.py
+++ b/cachalot/tests/write.py
@@ -9,9 +9,9 @@ except ImportError:  # For Python 2.6
 from django import VERSION as django_version
 from django.contrib.auth.models import User, Permission, Group
 from django.core.exceptions import MultipleObjectsReturned
-from django.db import connection
+from django.db import connection, transaction
 from django.db.models import Count
-from django.test import TransactionTestCase
+from django.test import TransactionTestCase, skipUnlessDBFeature
 
 from .models import Test
 
@@ -481,6 +481,27 @@ class WriteTestCase(TransactionTestCase):
                          .prefetch_related('owner__groups__permissions'))
             self.assertEqual(data7[0].owner.username, 'modified_user')
 
+    @skipUnlessDBFeature('has_select_for_update')
+    def test_invalidate_select_for_update(self):
+        with self.assertNumQueries(1):
+            Test.objects.bulk_create([Test(name='test1'), Test(name='test2')])
+
+        with self.assertNumQueries(1):
+            with transaction.atomic():
+                data1 = list(Test.objects.select_for_update())
+                self.assertListEqual([t.name for t in data1],
+                                     ['test1', 'test2'])
+
+        with self.assertNumQueries(1):
+            with transaction.atomic():
+                qs = Test.objects.select_for_update()
+                qs.update(name='test3')
+
+        with self.assertNumQueries(1):
+            with transaction.atomic():
+                data2 = list(Test.objects.select_for_update())
+                self.assertListEqual([t.name for t in data2], ['test3'] * 2)
+
     @skip(NotImplementedError)
     def test_invalidate_extra_select(self):
         pass