/tests/regressiontests/file_uploads/uploadhandler.py

 1"""
 2Upload handlers to test the upload API.
 3"""
 4
 5from django.core.files.uploadhandler import FileUploadHandler, StopUpload
 6
 7class QuotaUploadHandler(FileUploadHandler):
 8    """
 9    This test upload handler terminates the connection if more than a quota
10    (5MB) is uploaded.
11    """
12    
13    QUOTA = 5 * 2**20 # 5 MB
14    
15    def __init__(self, request=None):
16        super(QuotaUploadHandler, self).__init__(request)
17        self.total_upload = 0
18        
19    def receive_data_chunk(self, raw_data, start):
20        self.total_upload += len(raw_data)
21        if self.total_upload >= self.QUOTA:
22            raise StopUpload(connection_reset=True)
23        return raw_data
24            
25    def file_complete(self, file_size):
26        return None
27
28class CustomUploadError(Exception):
29    pass
30
31class ErroringUploadHandler(FileUploadHandler):
32    """A handler that raises an exception."""
33    def receive_data_chunk(self, raw_data, start):
34        raise CustomUploadError("Oops!")