diff --git a/modules/head.py b/modules/head.py index 235db58..8a5f55c 100644 --- a/modules/head.py +++ b/modules/head.py @@ -47,7 +47,10 @@ def head(phenny, input): info = web.head(uri) info['status'] = '200' except web.HTTPError as e: - return phenny.say(str(e.code)) + if hasattr(e, 'code'): + return phenny.say(str(e.code)) + else: + return phenny.say(str(e.response.status_code)) except web.ConnectionError: return phenny.say("Can't connect to %s" % uri) diff --git a/modules/test/test_head.py b/modules/test/test_head.py index fdfe95c..d92ceee 100644 --- a/modules/test/test_head.py +++ b/modules/test/test_head.py @@ -21,6 +21,13 @@ class TestHead(unittest.TestCase): '\d{2}:\d{2}:\d{2} UTC, [0-9\.]+ s$', out, flags=re.UNICODE) self.assertTrue(m) + def test_head_404(self): + input = Mock(group=lambda x: 'http://vtluug.org/trigger_404') + head(self.phenny, input) + + out = self.phenny.say.call_args[0][0] + self.assertEqual(out, '404') + def test_header(self): input = Mock(group=lambda x: 'http://vtluug.org Server') head(self.phenny, input)