Fix for bug that allowed taking over others’ chunks

src/shop/api.py

@@ -205,7 +205,7 @@ def finalize_stocktaking(stocktake_id):
 
 
 def finalize_stocktake_chunk(chunk_id):
-    """Markes given chunk as finished."""
+    """Marks given chunk as finished."""
     chunk_obj = models.StocktakeChunk.objects.get(id=chunk_id)
     if chunk_obj.locked:
         raise exceptions.APIException('Chunk already locked.')
@@ -216,12 +216,23 @@ def finalize_stocktake_chunk(chunk_id):
 
 @transaction.atomic
 def assign_free_stocktake_chunk(user_id, stocktake_id):
-    """Assignes a free stock-take chunk to the user, if any free left.
+    """Assigns a free stock-take chunk to a user, if any free left.
 
     If user is already assigned to a chunk, that chunk should be returned.
     """
     chunk_qs = models.StocktakeChunk.objects.select_for_update()
-    chunk_objs = chunk_qs.filter(stocktake_id=stocktake_id, locked=False)
+    try:
+        return chunk_qs.get(
+            stocktake_id=stocktake_id,
+            owner_id=user_id
+        )
+    except models.StocktakeChunk.DoesNotExist:
+        pass
+    chunk_objs = chunk_qs.filter(
+        stocktake_id=stocktake_id,
+        locked=False,
+        owner__is_null=True
+    )
     if not chunk_objs:
         return None
     chunk_obj = chunk_objs.first()

src/shop/migrations/0016_auto_20170214_0743.py

@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.10.5 on 2017-02-14 07:43
+from __future__ import unicode_literals
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('shop', '0015_auto_20170213_2104'),
+    ]
+
+    operations = [
+        migrations.AlterUniqueTogether(
+            name='stocktakechunk',
+            unique_together=set([('stocktake', 'locked')]),
+        ),
+    ]

src/shop/models.py

@@ -126,12 +126,13 @@ class StocktakeChunk(UUIDModel):
     locked = models.BooleanField(default=False)
     owner = models.ForeignKey('auth.User', null=True, blank=True)
 
-    def item_count(self):
-        return self.items.count()
-
     class Meta:
         verbose_name = _('Chunk')
         verbose_name_plural = _('Chunks')
+        unique_together = ('stocktake', 'locked')
+
+    def item_count(self):
+        return self.items.count()
 
     def __str__(self):
         return str(self.id)

src/shop/tests/test_api.py

@@ -290,21 +290,33 @@ def test_assign_free_stocktake_chunk(self):
             stocktake=stocktake_obj,
             locked=False
         )
-        user_obj = User.objects.create_superuser(
+        factories.StocktakeChunkFactory.create(
+            stocktake=stocktake_obj,
+            locked=False
+        )
+        user_obj1 = User.objects.create_superuser(
             'the_baconator', 'bacon@foobar.com', '123'
         )
-        obj1 = api.assign_free_stocktake_chunk(user_obj.id, stocktake_obj.id)
+        user_obj2 = User.objects.create_superuser(
+            'ludo', 'ludo@foobar.com', '123'
+        )
+        obj1 = api.assign_free_stocktake_chunk(user_obj1.id, stocktake_obj.id)
         self.assertIsNotNone(obj1)
         # Work on the first chunk not yet completed
-        obj2 = api.assign_free_stocktake_chunk(user_obj.id, stocktake_obj.id)
+        obj2 = api.assign_free_stocktake_chunk(user_obj1.id, stocktake_obj.id)
         self.assertIsNotNone(obj1)
         self.assertEqual(obj1.id, obj2.id)
         api.finalize_stocktake_chunk(obj1.id)
         # Work on the next chunk
-        obj3 = api.assign_free_stocktake_chunk(user_obj.id, stocktake_obj.id)
+        obj3 = api.assign_free_stocktake_chunk(user_obj1.id, stocktake_obj.id)
         self.assertIsNotNone(obj3)
         self.assertNotEqual(obj2.id, obj3.id)
+        # Another user wanst to work too
+        obj4 = api.assign_free_stocktake_chunk(user_obj2.id, stocktake_obj.id)
+        self.assertIsNotNone(obj4)
+        self.assertNotEqual(obj3.id, obj4.id)
         api.finalize_stocktake_chunk(obj3.id)
+        api.finalize_stocktake_chunk(obj4.id)
         # Work is finished
-        obj4 = api.assign_free_stocktake_chunk(user_obj.id, stocktake_obj.id)
-        self.assertIsNone(obj4)
+        obj5 = api.assign_free_stocktake_chunk(user_obj1.id, stocktake_obj.id)
+        self.assertIsNone(obj5)