github/webssh
1
0
Fork 0
mirror of https://github.com/huashengdun/webssh.git synced 2026-02-14 11:49:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added attribute custom_headers to MixinHandler

Browse Source
This commit is contained in:
Sheng
2018-10-16 15:14:34 +08:00
parent af60cd1cd5
commit a68eff592f
2 changed files with 17 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/test_app.py
View File

@@ -586,22 +586,30 @@ class TestAppWithTrustedStream(OtherTestBase):
class TestAppNotFoundHandler(OtherTestBase):
custom_headers = handler.MixinHandler.custom_headers
def test_with_not_found_get_request(self):
response = self.fetch('/pathnotfound', method='GET')
self.assertEqual(response.code, 404)
self.assertEqual(response.headers['Server'], 'TornadoServer')
self.assertEqual(
response.headers['Server'], self.custom_headers['Server']
)
self.assertIn(b'404: Not Found', response.body)
def test_with_not_found_post_request(self):
response = self.fetch('/pathnotfound', method='POST',
body=urlencode(self.body), headers=self.headers)
self.assertEqual(response.code, 404)
self.assertEqual(response.headers['Server'], 'TornadoServer')
self.assertEqual(
response.headers['Server'], self.custom_headers['Server']
)
self.assertIn(b'404: Not Found', response.body)
def test_with_not_found_put_request(self):
response = self.fetch('/pathnotfound', method='PUT',
body=urlencode(self.body), headers=self.headers)
self.assertEqual(response.code, 404)
self.assertEqual(response.headers['Server'], 'TornadoServer')
self.assertEqual(
response.headers['Server'], self.custom_headers['Server']
)
self.assertIn(b'404: Not Found', response.body)

7
webssh/handler.py
View File

@@ -39,6 +39,10 @@ class InvalidValueError(Exception):
class MixinHandler(object):
custom_headers = {
'Server': 'TornadoServer'
}
def prepare(self):
if self.is_forbidden():
raise tornado.web.HTTPError(403)
@@ -66,7 +70,8 @@ class MixinHandler(object):
return True
def set_default_headers(self):
self.set_header('Server', 'TornadoServer')
for header in self.custom_headers.items():
self.set_header(*header)
def get_value(self, name):
value = self.get_argument(name)