|
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