Ticket #19094: 19094-1.diff

File 19094-1.diff, 6.3 KB (added by claudep, 2 years ago)

Improved FakePayload to support write, len and string input

  • django/test/client.py

    diff --git a/django/test/client.py b/django/test/client.py
    index 8fd765e..6d12321 100644
    a b class FakePayload(object): 
    4343    length. This makes sure that views can't do anything under the test client
    4444    that wouldn't work in Real Life.
    4545    """
    46     def __init__(self, content):
    47         self.__content = BytesIO(content)
    48         self.__len = len(content)
     46    def __init__(self, content=None):
     47        self.__content = BytesIO()
     48        self.__len = 0
     49        self.read_started = False
     50        if content is not None:
     51            self.write(content)
     52
     53    def __len__(self):
     54        return self.__len
    4955
    5056    def read(self, num_bytes=None):
     57        if not self.read_started:
     58            self.__content.seek(0)
     59            self.read_started = True
    5160        if num_bytes is None:
    5261            num_bytes = self.__len or 0
    5362        assert self.__len >= num_bytes, "Cannot read more than the available bytes from the HTTP incoming data."
    class FakePayload(object): 
    5564        self.__len -= num_bytes
    5665        return content
    5766
     67    def write(self, content):
     68        if self.read_started:
     69            raise ValueError("Unable to write a payload after he's been read")
     70        content = force_bytes(content)
     71        self.__content.write(content)
     72        self.__len += len(content)
     73
    5874
    5975class ClientHandler(BaseHandler):
    6076    """
  • tests/regressiontests/file_uploads/tests.py

    diff --git a/tests/regressiontests/file_uploads/tests.py b/tests/regressiontests/file_uploads/tests.py
    index 2a1ec7d..918f77d 100644
    a b class FileUploadTests(TestCase): 
    6262
    6363    def test_base64_upload(self):
    6464        test_string = "This data will be transmitted base64-encoded."
    65         payload = "\r\n".join([
     65        payload = client.FakePayload("\r\n".join([
    6666            '--' + client.BOUNDARY,
    6767            'Content-Disposition: form-data; name="file"; filename="test.txt"',
    6868            'Content-Type: application/octet-stream',
    6969            'Content-Transfer-Encoding: base64',
    70             '',
    71             base64.b64encode(force_bytes(test_string)).decode('ascii'),
    72             '--' + client.BOUNDARY + '--',
    73             '',
    74         ]).encode('utf-8')
     70            '',]))
     71        payload.write(b"\r\n" + base64.b64encode(force_bytes(test_string)) + b"\r\n")
     72        payload.write('--' + client.BOUNDARY + '--\r\n')
    7573        r = {
    7674            'CONTENT_LENGTH': len(payload),
    7775            'CONTENT_TYPE':   client.MULTIPART_CONTENT,
    7876            'PATH_INFO':      "/file_uploads/echo_content/",
    7977            'REQUEST_METHOD': 'POST',
    80             'wsgi.input':     client.FakePayload(payload),
     78            'wsgi.input':     payload,
    8179        }
    8280        response = self.client.request(**r)
    8381        received = json.loads(response.content.decode('utf-8'))
    class FileUploadTests(TestCase): 
    126124            "../..\\hax0rd.txt"         # Relative path, mixed.
    127125        ]
    128126
    129         payload = []
     127        payload = client.FakePayload()
    130128        for i, name in enumerate(scary_file_names):
    131             payload.extend([
     129            payload.write('\r\n'.join([
    132130                '--' + client.BOUNDARY,
    133131                'Content-Disposition: form-data; name="file%s"; filename="%s"' % (i, name),
    134132                'Content-Type: application/octet-stream',
    135133                '',
    136                 'You got pwnd.'
    137             ])
    138         payload.extend([
    139             '--' + client.BOUNDARY + '--',
    140             '',
    141         ])
     134                'You got pwnd.\r\n'
     135            ]))
     136        payload.write('\r\n--' + client.BOUNDARY + '--\r\n')
    142137
    143         payload = "\r\n".join(payload).encode('utf-8')
    144138        r = {
    145139            'CONTENT_LENGTH': len(payload),
    146140            'CONTENT_TYPE':   client.MULTIPART_CONTENT,
    147141            'PATH_INFO':      "/file_uploads/echo/",
    148142            'REQUEST_METHOD': 'POST',
    149             'wsgi.input':     client.FakePayload(payload),
     143            'wsgi.input':     payload,
    150144        }
    151145        response = self.client.request(**r)
    152146
    class FileUploadTests(TestCase): 
    159153    def test_filename_overflow(self):
    160154        """File names over 256 characters (dangerous on some platforms) get fixed up."""
    161155        name = "%s.txt" % ("f"*500)
    162         payload = "\r\n".join([
     156        payload = client.FakePayload("\r\n".join([
    163157            '--' + client.BOUNDARY,
    164158            'Content-Disposition: form-data; name="file"; filename="%s"' % name,
    165159            'Content-Type: application/octet-stream',
    class FileUploadTests(TestCase): 
    167161            'Oops.'
    168162            '--' + client.BOUNDARY + '--',
    169163            '',
    170         ]).encode('utf-8')
     164        ]))
    171165        r = {
    172166            'CONTENT_LENGTH': len(payload),
    173167            'CONTENT_TYPE':   client.MULTIPART_CONTENT,
    174168            'PATH_INFO':      "/file_uploads/echo/",
    175169            'REQUEST_METHOD': 'POST',
    176             'wsgi.input':     client.FakePayload(payload),
     170            'wsgi.input':     payload,
    177171        }
    178172        got = json.loads(self.client.request(**r).content.decode('utf-8'))
    179173        self.assertTrue(len(got['file']) < 256, "Got a long file name (%s characters)." % len(got['file']))
    class FileUploadTests(TestCase): 
    184178        attempt to read beyond the end of the stream, and simply will handle
    185179        the part that can be parsed gracefully.
    186180        """
    187         payload = "\r\n".join([
     181        payload_str = "\r\n".join([
    188182            '--' + client.BOUNDARY,
    189183            'Content-Disposition: form-data; name="file"; filename="foo.txt"',
    190184            'Content-Type: application/octet-stream',
    class FileUploadTests(TestCase): 
    192186            'file contents'
    193187            '--' + client.BOUNDARY + '--',
    194188            '',
    195         ]).encode('utf-8')
    196         payload = payload[:-10]
     189        ])
     190        payload = client.FakePayload(payload_str[:-10])
    197191        r = {
    198192            'CONTENT_LENGTH': len(payload),
    199193            'CONTENT_TYPE': client.MULTIPART_CONTENT,
    200194            'PATH_INFO': '/file_uploads/echo/',
    201195            'REQUEST_METHOD': 'POST',
    202             'wsgi.input': client.FakePayload(payload),
     196            'wsgi.input': payload,
    203197        }
    204198        got = json.loads(self.client.request(**r).content.decode('utf-8'))
    205199        self.assertEqual(got, {})
Back to Top