Adds set_quantity as a method on CartController.

Refactors add_to_cart to be in terms of set_quantity

registrasion/controllers/cart.py

@@ -54,40 +54,73 @@ class CartController(object):
         self.cart.time_last_updated = timezone.now()
         self.cart.reservation_duration = max(reservations)
 
-    def add_to_cart(self, product, quantity):
-        ''' Adds _quantity_ of the given _product_ to the cart. Raises
-        ValidationError if constraints are violated.'''
-
-        prod = ProductController(product)
-
-        # TODO: Check enabling conditions for product for user
-
-        if not prod.can_add_with_enabling_conditions(self.cart.user, quantity):
-            raise ValidationError("Not enough of that product left (ec)")
-
-        if not prod.user_can_add_within_limit(self.cart.user, quantity):
-            raise ValidationError("Not enough of that product left (user)")
-
-        try:
-            # Try to update an existing item within this cart if possible.
-            product_item = rego.ProductItem.objects.get(
-                cart=self.cart,
-                product=product)
-            product_item.quantity += quantity
-        except ObjectDoesNotExist:
-            product_item = rego.ProductItem.objects.create(
-                cart=self.cart,
-                product=product,
-                quantity=quantity,
-            )
-        product_item.save()
-
+    def end_batch(self):
+        ''' Performs operations that occur occur at the end of a batch of
+        product changes/voucher applications etc. '''
         self.recalculate_discounts()
 
         self.extend_reservation()
         self.cart.revision += 1
         self.cart.save()
 
+    def set_quantity(self, product, quantity, batched=False):
+        ''' Sets the _quantity_ of the given _product_ in the cart to the given
+        _quantity_. '''
+
+        if quantity < 0:
+            raise ValidationError("Cannot have fewer than 0 items in cart.")
+
+        try:
+            product_item = rego.ProductItem.objects.get(
+                cart=self.cart,
+                product=product)
+            old_quantity = product_item.quantity
+
+            if quantity == 0:
+                product_item.delete()
+                return
+        except ObjectDoesNotExist:
+            if quantity == 0:
+                return
+
+            product_item = rego.ProductItem.objects.create(
+                cart=self.cart,
+                product=product,
+                quantity=0,
+            )
+
+            old_quantity = 0
+
+        # Validate the addition to the cart
+        adjustment = quantity - old_quantity
+        prod = ProductController(product)
+
+        if not prod.can_add_with_enabling_conditions(
+                self.cart.user, adjustment):
+            raise ValidationError("Not enough of that product left (ec)")
+
+        if not prod.user_can_add_within_limit(self.cart.user, adjustment):
+            raise ValidationError("Not enough of that product left (user)")
+
+        product_item.quantity = quantity
+        product_item.save()
+
+        if not batched:
+            self.end_batch()
+
+    def add_to_cart(self, product, quantity):
+        ''' Adds _quantity_ of the given _product_ to the cart. Raises
+        ValidationError if constraints are violated.'''
+
+        try:
+            product_item = rego.ProductItem.objects.get(
+                cart=self.cart,
+                product=product)
+            old_quantity = product_item.quantity
+        except ObjectDoesNotExist:
+            old_quantity = 0
+        self.set_quantity(product, old_quantity + quantity)
+
     def apply_voucher(self, voucher):
         ''' Applies the given voucher to this cart. '''
 
@@ -101,10 +134,7 @@ class CartController(object):
 
         # If successful...
         self.cart.vouchers.add(voucher)
-
-        self.extend_reservation()
-        self.cart.revision += 1
-        self.cart.save()
+        self.end_batch()
 
     def validate_cart(self):
         ''' Determines whether the status of the current cart is valid;

registrasion/tests/test_cart.py

@@ -3,6 +3,7 @@ import pytz
 
 from decimal import Decimal
 from django.contrib.auth.models import User
+from django.core.exceptions import ObjectDoesNotExist
 from django.core.exceptions import ValidationError
 from django.test import TestCase
 
@@ -159,6 +160,38 @@ class BasicCartTests(RegistrationCartTestCase):
         item = items[0]
         self.assertEquals(2, item.quantity)
 
+    def test_set_quantity(self):
+        current_cart = CartController.for_user(self.USER_1)
+
+        def get_item():
+            return rego.ProductItem.objects.get(
+                cart=current_cart.cart,
+                product=self.PROD_1)
+
+        current_cart.set_quantity(self.PROD_1, 1)
+        self.assertEqual(1, get_item().quantity)
+
+        # Setting the quantity to zero should remove the entry from the cart.
+        current_cart.set_quantity(self.PROD_1, 0)
+        with self.assertRaises(ObjectDoesNotExist):
+            get_item()
+
+        current_cart.set_quantity(self.PROD_1, 9)
+        self.assertEqual(9, get_item().quantity)
+
+        with self.assertRaises(ValidationError):
+            current_cart.set_quantity(self.PROD_1, 11)
+
+        self.assertEqual(9, get_item().quantity)
+
+        with self.assertRaises(ValidationError):
+            current_cart.set_quantity(self.PROD_1, -1)
+
+        self.assertEqual(9, get_item().quantity)
+
+        current_cart.set_quantity(self.PROD_1, 2)
+        self.assertEqual(2, get_item().quantity)
+
     def test_add_to_cart_per_user_limit(self):
         current_cart = CartController.for_user(self.USER_1)