|
15 | 15 | """Tests for nova websocketproxy."""
|
16 | 16 |
|
17 | 17 | import copy
|
| 18 | +import io |
18 | 19 | import socket
|
19 | 20 |
|
20 | 21 | import mock
|
@@ -635,30 +636,28 @@ def test_reject_open_redirect(self):
|
635 | 636 | b''
|
636 | 637 | ]
|
637 | 638 |
|
638 |
| - # Collect the response data to verify at the end. The |
639 |
| - # SimpleHTTPRequestHandler writes the response data by calling the |
640 |
| - # request socket sendall() method. |
641 |
| - self.data = b'' |
642 |
| - |
643 |
| - def fake_sendall(data): |
644 |
| - self.data += data |
645 |
| - |
646 |
| - mock_req.sendall.side_effect = fake_sendall |
647 |
| - |
648 | 639 | client_addr = ('8.8.8.8', 54321)
|
649 | 640 | mock_server = mock.MagicMock()
|
650 | 641 | # This specifies that the server will be able to handle requests other
|
651 | 642 | # than only websockets.
|
652 | 643 | mock_server.only_upgrade = False
|
653 | 644 |
|
654 | 645 | # Constructing a handler will process the mock_req request passed in.
|
655 |
| - websocketproxy.NovaProxyRequestHandler( |
| 646 | + handler = websocketproxy.NovaProxyRequestHandler( |
656 | 647 | mock_req, client_addr, mock_server)
|
657 | 648 |
|
| 649 | + # Collect the response data to verify at the end. The |
| 650 | + # SimpleHTTPRequestHandler writes the response data to a 'wfile' |
| 651 | + # attribute. |
| 652 | + output = io.BytesIO() |
| 653 | + handler.wfile = output |
| 654 | + # Process the mock_req again to do the capture. |
| 655 | + handler.do_GET() |
| 656 | + output.seek(0) |
| 657 | + result = output.readlines() |
| 658 | + |
658 | 659 | # Verify no redirect happens and instead a 400 Bad Request is returned.
|
659 |
| - self.data = self.data.decode() |
660 |
| - self.assertIn('Error code: 400', self.data) |
661 |
| - self.assertIn('Message: URI must not start with //', self.data) |
| 660 | + self.assertIn('400 URI must not start with //', result[0].decode()) |
662 | 661 |
|
663 | 662 | @mock.patch('websockify.websocketproxy.select_ssl_version')
|
664 | 663 | def test_ssl_min_version_is_not_set(self, mock_select_ssl):
|
|
0 commit comments