webssh

commit 50ab045c1f1678e3758a8578f3aae8178a3f124b
parent 42ac3daf87d5c4cf2961c147dca8ab2fd1a6d9e7
Author: Sheng <webmaster0115@gmail.com>
Date:   Tue, 24 Apr 2018 10:57:25 +0800

Use HTTPServerRequest instead

Diffstat:
M
tests/test_handler.py
 | 
27
++++++++-------------------

1 file changed, 8 insertions(+), 19 deletions(-)
diff --git a/tests/test_handler.py b/tests/test_handler.py
@@ -1,50 +1,39 @@
 import unittest
 import sys
 
+from tornado.httputil import HTTPServerRequest
 from handler import MixinHandler
 
 
-class RequestMock(object):
-
-    def __init__(self):
-        self.headers = {}
-
-    def set_ip(self, ip):
-        self.headers['X-Real-Ip'] = ip
-
-    def set_port(self, port):
-        self.headers['X-Real-Port'] = port
-
-
 class TestMixinHandler(unittest.TestCase):
 
     def test_get_real_client_addr_without_nginx_config(self):
         handler = MixinHandler()
-        handler.request = RequestMock()
+        handler.request = HTTPServerRequest(uri='/')
         self.assertIsNone(handler.get_real_client_addr())
 
     def test_get_real_client_addr_with_correct_nginx_config(self):
         handler = MixinHandler()
-        handler.request = RequestMock()
+        handler.request = HTTPServerRequest(uri='/')
 
         ip = '127.0.0.1'
-        handler.request.set_ip(ip)
-        handler.request.set_port('12345')
+        handler.request.headers.add('X-Real-Ip', ip)
+        handler.request.headers.add('X-Real-Port', '12345')
         self.assertEqual(handler.get_real_client_addr(), (ip, 12345))
 
     @unittest.skipIf(sys.version_info < (3,),
                      reason='assertLogs not supported in Python 2')
     def test_get_real_client_addr_with_bad_nginx_config(self):
         handler = MixinHandler()
-        handler.request = RequestMock()
+        handler.request = HTTPServerRequest(uri='/')
 
         ip = '127.0.0.1'
-        handler.request.set_ip(ip)
+        handler.request.headers.add('X-Real-Ip', ip)
         with self.assertLogs() as cm:
             handler.get_real_client_addr()
         self.assertEqual(cm.output, ['WARNING:root:Bad nginx configuration.'])
 
-        handler.request.set_port('12345x')
+        handler.request.headers.add('X-Real-Port', '12345x')
         with self.assertLogs() as cm:
             handler.get_real_client_addr()
         self.assertEqual(cm.output, ['WARNING:root:Bad nginx configuration.'])