Add quota validator

Author: hirunatan

validatedfile/models.py

@@ -32,3 +32,31 @@ def clean(self, *args, **kwargs):
 
         return data
 
+
+class FileQuota(object):
+    def __init__(self, items = [], file_attr_name = None):
+        self.items = items
+        self.file_attr_name = file_attr_name
+
+    def current_usage(self):
+        usage = 0
+        for item in self.items:
+            the_file = getattr(item, self.file_attr_name, None)
+            if the_file:
+                usage += the_file.size
+        return usage
+
+class QuotaValidator(object):
+    def __init__(self, max_usage):
+        self.quota = FileQuota()
+        self.max_usage = max_usage
+
+    def set_quota(self, quota):
+        self.quota = quota
+
+    def __call__(self, file):
+        tried_usage = self.quota.current_usage() + file.size
+        if tried_usage > self.max_usage:
+            raise forms.ValidationError(_('Please keep the total uploaded files under %s. With this file, the total would be %s' %
+                                   (filesizeformat(self.max_usage), filesizeformat(tried_usage))))
+        

validatedfile/tests/__init__.py

@@ -5,8 +5,10 @@
 
 import os.path
 
-from models import TestModel, TestModelNoValidate
-from forms import TestModelForm, TestModelNoValidateForm
+from validatedfile.models import FileQuota
+
+from models import TestModel, TestModelNoValidate, TestContainer, TestElement
+from forms import TestModelForm, TestModelNoValidateForm, TestElementForm
 
 class ValidatedFileFieldTest(TestCase):
 
@@ -15,14 +17,12 @@ class ValidatedFileFieldTest(TestCase):
 
     def test_create_empty_instance(self):
         instance = TestModel.objects.create()
-        instance.save()
 
 
     def test_create_instance_with_file(self):
         instance = TestModel.objects.create(
                 the_file = File(self._get_sample_file('image2k.png'), 'the_file.png')
             )
-        instance.save()
 
         self._check_file_url(instance.the_file, 'the_file.png')
 
@@ -120,6 +120,107 @@ def test_form_null_file(self):
         instance.delete()
 
 
+    def test_quota_empty(self):
+        container = TestContainer.objects.create(name = 'container1')
+
+        quota = FileQuota(container.test_elements.all(), 'the_file')
+        self.assertEqual(quota.current_usage(), 0)
+
+        container.delete()
+
+
+    def test_quota_one_file(self):
+        container = TestContainer.objects.create(name = 'container1')
+        element = TestElement.objects.create(
+                container = container,
+                the_file = File(self._get_sample_file('image2k.png'), 'the_file.png')
+            )
+
+        quota = FileQuota(container.test_elements.all(), 'the_file')
+        self.assertEqual(quota.current_usage(), 2120)
+
+        element.the_file.delete()
+        element.delete()
+        container.delete()
+
+
+    def test_quota_several_files_several_containers(self):
+        container1 = TestContainer.objects.create(name = 'container1')
+        element1 = TestElement.objects.create(
+                container = container1,
+                the_file = File(self._get_sample_file('image2k.png'), 'the_file1.png')
+            )
+        element2 = TestElement.objects.create(
+                container = container1,
+                the_file = File(self._get_sample_file('image15k.png'), 'the_file2.png')
+            )
+        container2 = TestContainer.objects.create(name = 'container2')
+        element3 = TestElement.objects.create(
+                container = container2,
+                the_file = File(self._get_sample_file('document15k.pdf'), 'the_file3.pdf')
+            )
+
+        quota = FileQuota(container1.test_elements.all(), 'the_file')
+        self.assertEqual(quota.current_usage(), 16706)
+
+        element1.the_file.delete()
+        element2.the_file.delete()
+        element3.the_file.delete()
+        element1.delete()
+        element2.delete()
+        element3.delete()
+        container1.delete()
+        container2.delete()
+
+
+    def test_form_quota_ok(self):
+        container = TestContainer.objects.create(name = 'container1')
+
+        quota = FileQuota(container.test_elements.all(), 'the_file')
+        uploaded_file = SimpleUploadedFile(
+                name = 'the_file.png',
+                content = self._get_sample_file('image2k.png').read(),
+                content_type = 'image/png',
+            )
+        form = TestElementForm(
+                container = container,
+                quota = quota,
+                data = {},
+                files = {'the_file': uploaded_file}
+            )
+
+        self.assertTrue(form.is_valid())
+
+        container.delete()
+
+
+    def test_form_quota_exceeded(self):
+        container = TestContainer.objects.create(name = 'container1')
+        element = TestElement.objects.create(
+                container = container,
+                the_file = File(self._get_sample_file('image2k.png'), 'the_file.png')
+            )
+
+        quota = FileQuota(container.test_elements.all(), 'the_file')
+        uploaded_file = SimpleUploadedFile(
+                name = 'the_file.png',
+                content = self._get_sample_file('image15k.png').read(),
+                content_type = 'image/png',
+            )
+        form = TestElementForm(
+                container = container,
+                quota = quota,
+                data = {},
+                files = {'the_file': uploaded_file}
+            )
+
+        self.assertFalse(form.is_valid())
+
+        element.the_file.delete()
+        element.delete()
+        container.delete()
+
+
     # Utilities
 
     def _get_sample_file(self, filename):

validatedfile/tests/forms.py

@@ -1,13 +1,35 @@
 from django import forms
-from models import TestModel, TestModelNoValidate
+from models import TestModel, TestModelNoValidate, TestContainer, TestElement
+from validatedfile.models import QuotaValidator
 
 class TestModelForm(forms.ModelForm):
 
     class Meta:
         model = TestModel
 
+
 class TestModelNoValidateForm(forms.ModelForm):
 
     class Meta:
         model = TestModelNoValidate
 
+
+class TestElementForm(forms.ModelForm):
+
+    the_file = forms.FileField(required = False,
+                    validators = [QuotaValidator(max_usage = 10000)])
+
+    class Meta:
+        model = TestElement
+        fields = ['the_file']
+
+    def __init__(self, container, quota, *args, **kwargs):
+        super(TestElementForm, self).__init__(*args, **kwargs)
+        self.container = container
+        self.fields['the_file'].validators[0].set_quota(quota)
+
+    def save(self, *args, **kwargs):
+        element = super(TestElementForm, self).save(commit = False)
+        element.container = self.container
+        element.save()
+

validatedfile/tests/models.py

@@ -1,5 +1,5 @@
 from django.db import models
-from validatedfile.models import ValidatedFileField
+from validatedfile.models import ValidatedFileField, QuotaValidator
 
 class TestModel(models.Model):
     the_file = ValidatedFileField(
@@ -15,3 +15,18 @@ class TestModelNoValidate(models.Model):
                     blank = True,
                     upload_to = 'testfile')
 
+
+class TestContainer(models.Model):
+    name = models.CharField(max_length = 100)
+
+
+class TestElement(models.Model):
+    container = models.ForeignKey(
+                    TestContainer,
+                    related_name = 'test_elements')
+    the_file = ValidatedFileField(
+                    null = True,
+                    blank = True,
+                    upload_to = 'testfile',
+                    content_types = ['image/png', 'image/jpeg'])
+