parts/django/tests/regressiontests/file_uploads/uploadhandler.py
changeset 307 c6bca38c1cbf
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/parts/django/tests/regressiontests/file_uploads/uploadhandler.py	Sat Jan 08 11:20:57 2011 +0530
@@ -0,0 +1,34 @@
+"""
+Upload handlers to test the upload API.
+"""
+
+from django.core.files.uploadhandler import FileUploadHandler, StopUpload
+
+class QuotaUploadHandler(FileUploadHandler):
+    """
+    This test upload handler terminates the connection if more than a quota
+    (5MB) is uploaded.
+    """
+    
+    QUOTA = 5 * 2**20 # 5 MB
+    
+    def __init__(self, request=None):
+        super(QuotaUploadHandler, self).__init__(request)
+        self.total_upload = 0
+        
+    def receive_data_chunk(self, raw_data, start):
+        self.total_upload += len(raw_data)
+        if self.total_upload >= self.QUOTA:
+            raise StopUpload(connection_reset=True)
+        return raw_data
+            
+    def file_complete(self, file_size):
+        return None
+
+class CustomUploadError(Exception):
+    pass
+
+class ErroringUploadHandler(FileUploadHandler):
+    """A handler that raises an exception."""
+    def receive_data_chunk(self, raw_data, start):
+        raise CustomUploadError("Oops!")