Code

Ticket #19094: 19094-1.diff

File 19094-1.diff, 6.3 KB (added by claudep, 21 months ago)

Improved FakePayload to support write, len and string input

Line 
1diff --git a/django/test/client.py b/django/test/client.py
2index 8fd765e..6d12321 100644
3--- a/django/test/client.py
4+++ b/django/test/client.py
5@@ -43,11 +43,20 @@ class FakePayload(object):
6     length. This makes sure that views can't do anything under the test client
7     that wouldn't work in Real Life.
8     """
9-    def __init__(self, content):
10-        self.__content = BytesIO(content)
11-        self.__len = len(content)
12+    def __init__(self, content=None):
13+        self.__content = BytesIO()
14+        self.__len = 0
15+        self.read_started = False
16+        if content is not None:
17+            self.write(content)
18+
19+    def __len__(self):
20+        return self.__len
21 
22     def read(self, num_bytes=None):
23+        if not self.read_started:
24+            self.__content.seek(0)
25+            self.read_started = True
26         if num_bytes is None:
27             num_bytes = self.__len or 0
28         assert self.__len >= num_bytes, "Cannot read more than the available bytes from the HTTP incoming data."
29@@ -55,6 +64,13 @@ class FakePayload(object):
30         self.__len -= num_bytes
31         return content
32 
33+    def write(self, content):
34+        if self.read_started:
35+            raise ValueError("Unable to write a payload after he's been read")
36+        content = force_bytes(content)
37+        self.__content.write(content)
38+        self.__len += len(content)
39+
40 
41 class ClientHandler(BaseHandler):
42     """
43diff --git a/tests/regressiontests/file_uploads/tests.py b/tests/regressiontests/file_uploads/tests.py
44index 2a1ec7d..918f77d 100644
45--- a/tests/regressiontests/file_uploads/tests.py
46+++ b/tests/regressiontests/file_uploads/tests.py
47@@ -62,22 +62,20 @@ class FileUploadTests(TestCase):
48 
49     def test_base64_upload(self):
50         test_string = "This data will be transmitted base64-encoded."
51-        payload = "\r\n".join([
52+        payload = client.FakePayload("\r\n".join([
53             '--' + client.BOUNDARY,
54             'Content-Disposition: form-data; name="file"; filename="test.txt"',
55             'Content-Type: application/octet-stream',
56             'Content-Transfer-Encoding: base64',
57-            '',
58-            base64.b64encode(force_bytes(test_string)).decode('ascii'),
59-            '--' + client.BOUNDARY + '--',
60-            '',
61-        ]).encode('utf-8')
62+            '',]))
63+        payload.write(b"\r\n" + base64.b64encode(force_bytes(test_string)) + b"\r\n")
64+        payload.write('--' + client.BOUNDARY + '--\r\n')
65         r = {
66             'CONTENT_LENGTH': len(payload),
67             'CONTENT_TYPE':   client.MULTIPART_CONTENT,
68             'PATH_INFO':      "/file_uploads/echo_content/",
69             'REQUEST_METHOD': 'POST',
70-            'wsgi.input':     client.FakePayload(payload),
71+            'wsgi.input':     payload,
72         }
73         response = self.client.request(**r)
74         received = json.loads(response.content.decode('utf-8'))
75@@ -126,27 +124,23 @@ class FileUploadTests(TestCase):
76             "../..\\hax0rd.txt"         # Relative path, mixed.
77         ]
78 
79-        payload = []
80+        payload = client.FakePayload()
81         for i, name in enumerate(scary_file_names):
82-            payload.extend([
83+            payload.write('\r\n'.join([
84                 '--' + client.BOUNDARY,
85                 'Content-Disposition: form-data; name="file%s"; filename="%s"' % (i, name),
86                 'Content-Type: application/octet-stream',
87                 '',
88-                'You got pwnd.'
89-            ])
90-        payload.extend([
91-            '--' + client.BOUNDARY + '--',
92-            '',
93-        ])
94+                'You got pwnd.\r\n'
95+            ]))
96+        payload.write('\r\n--' + client.BOUNDARY + '--\r\n')
97 
98-        payload = "\r\n".join(payload).encode('utf-8')
99         r = {
100             'CONTENT_LENGTH': len(payload),
101             'CONTENT_TYPE':   client.MULTIPART_CONTENT,
102             'PATH_INFO':      "/file_uploads/echo/",
103             'REQUEST_METHOD': 'POST',
104-            'wsgi.input':     client.FakePayload(payload),
105+            'wsgi.input':     payload,
106         }
107         response = self.client.request(**r)
108 
109@@ -159,7 +153,7 @@ class FileUploadTests(TestCase):
110     def test_filename_overflow(self):
111         """File names over 256 characters (dangerous on some platforms) get fixed up."""
112         name = "%s.txt" % ("f"*500)
113-        payload = "\r\n".join([
114+        payload = client.FakePayload("\r\n".join([
115             '--' + client.BOUNDARY,
116             'Content-Disposition: form-data; name="file"; filename="%s"' % name,
117             'Content-Type: application/octet-stream',
118@@ -167,13 +161,13 @@ class FileUploadTests(TestCase):
119             'Oops.'
120             '--' + client.BOUNDARY + '--',
121             '',
122-        ]).encode('utf-8')
123+        ]))
124         r = {
125             'CONTENT_LENGTH': len(payload),
126             'CONTENT_TYPE':   client.MULTIPART_CONTENT,
127             'PATH_INFO':      "/file_uploads/echo/",
128             'REQUEST_METHOD': 'POST',
129-            'wsgi.input':     client.FakePayload(payload),
130+            'wsgi.input':     payload,
131         }
132         got = json.loads(self.client.request(**r).content.decode('utf-8'))
133         self.assertTrue(len(got['file']) < 256, "Got a long file name (%s characters)." % len(got['file']))
134@@ -184,7 +178,7 @@ class FileUploadTests(TestCase):
135         attempt to read beyond the end of the stream, and simply will handle
136         the part that can be parsed gracefully.
137         """
138-        payload = "\r\n".join([
139+        payload_str = "\r\n".join([
140             '--' + client.BOUNDARY,
141             'Content-Disposition: form-data; name="file"; filename="foo.txt"',
142             'Content-Type: application/octet-stream',
143@@ -192,14 +186,14 @@ class FileUploadTests(TestCase):
144             'file contents'
145             '--' + client.BOUNDARY + '--',
146             '',
147-        ]).encode('utf-8')
148-        payload = payload[:-10]
149+        ])
150+        payload = client.FakePayload(payload_str[:-10])
151         r = {
152             'CONTENT_LENGTH': len(payload),
153             'CONTENT_TYPE': client.MULTIPART_CONTENT,
154             'PATH_INFO': '/file_uploads/echo/',
155             'REQUEST_METHOD': 'POST',
156-            'wsgi.input': client.FakePayload(payload),
157+            'wsgi.input': payload,
158         }
159         got = json.loads(self.client.request(**r).content.decode('utf-8'))
160         self.assertEqual(got, {})