parts/django/tests/regressiontests/file_uploads/uploadhandler.py
author Nishanth Amuluru <nishanth@fossee.in>
Sat, 08 Jan 2011 11:20:57 +0530
changeset 307 c6bca38c1cbf
permissions -rw-r--r--
Added buildout stuff and made changes accordingly
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
307
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
     1
"""
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
     2
Upload handlers to test the upload API.
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
     3
"""
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
     4
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
     5
from django.core.files.uploadhandler import FileUploadHandler, StopUpload
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
     6
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
     7
class QuotaUploadHandler(FileUploadHandler):
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
     8
    """
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
     9
    This test upload handler terminates the connection if more than a quota
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    10
    (5MB) is uploaded.
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    11
    """
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    12
    
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    13
    QUOTA = 5 * 2**20 # 5 MB
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    14
    
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    15
    def __init__(self, request=None):
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    16
        super(QuotaUploadHandler, self).__init__(request)
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    17
        self.total_upload = 0
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    18
        
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    19
    def receive_data_chunk(self, raw_data, start):
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    20
        self.total_upload += len(raw_data)
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    21
        if self.total_upload >= self.QUOTA:
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    22
            raise StopUpload(connection_reset=True)
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    23
        return raw_data
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    24
            
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    25
    def file_complete(self, file_size):
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    26
        return None
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    27
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    28
class CustomUploadError(Exception):
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    29
    pass
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    30
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    31
class ErroringUploadHandler(FileUploadHandler):
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    32
    """A handler that raises an exception."""
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    33
    def receive_data_chunk(self, raw_data, start):
c6bca38c1cbf Added buildout stuff and made changes accordingly
Nishanth Amuluru <nishanth@fossee.in>
parents:
diff changeset
    34
        raise CustomUploadError("Oops!")