redis support

hacheck/checker.py

@@ -145,3 +145,42 @@ def timed_out(duration):
         raise tornado.gen.Return((500, 'MySQL sez %s' % response))
     yield conn.quit()
     raise tornado.gen.Return((200, 'MySQL connect response: %s' % response))
+
+
+@cache.cached
+@tornado.gen.coroutine
+def check_redis(service_name, port, query, io_loop, query_params):
+    stream = None
+    connect_start = time.time()
+
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
+    try:
+        stream = tornado.iostream.IOStream(s, io_loop=io_loop)
+        yield add_timeout_to_task(
+            stream.connect,
+            args=[('127.0.0.1', port)],
+            timeout_secs=TIMEOUT
+        )
+        yield stream.write(b'PING\r\n')
+        response = yield add_timeout_to_task(
+            stream.read_until,
+            args=[b'\n'],
+            timeout_secs=1
+        )
+        stream.close()
+        if response.strip() != b'+PONG':
+            raise tornado.gen.Return((500, 'Sent PING, got back %s' % response))
+        else:
+            raise tornado.gen.Return((200, 'Sent PING, got back +PONG'))
+    except Timeout:
+        raise tornado.gen.Return((
+            503,
+            'Connection timed out after %.2fs' % (time.time() - connect_start)
+        ))
+    finally:
+        if stream:
+            stream.close()
+    raise tornado.gen.Return((
+        200,
+        'Connected in %.2fs' % (time.time() - connect_start)
+    ))

hacheck/handlers.py

@@ -83,3 +83,7 @@ class TCPServiceHandler(BaseServiceHandler):
 
 class MySQLServiceHandler(BaseServiceHandler):
     CHECKERS = [checker.check_spool, checker.check_mysql]
+
+
+class RedisServiceHandler(BaseServiceHandler):
+    CHECKERS = [checker.check_spool, checker.check_redis]

hacheck/main.py

@@ -18,6 +18,7 @@ def get_app():
         (r'/http/([a-zA-Z0-9_-]+)/([0-9]+)/(.*)', handlers.HTTPServiceHandler),
         (r'/tcp/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.TCPServiceHandler),
         (r'/mysql/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.MySQLServiceHandler),
+        (r'/redis/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.RedisServiceHandler),
         (r'/spool/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.SpoolServiceHandler),
         (r'/recent', handlers.ListRecentHandler),
         (r'/status', handlers.StatusHandler),

tests/test_checker.py

@@ -108,9 +108,13 @@ def test_query_params_not_passed(self):
 
 
 class TestServer(tornado.tcpserver.TCPServer):
+    def __init__(self, io_loop, response='hello\n'):
+        self.response = response
+        super(TestServer, self).__init__(io_loop=io_loop)
+
     @tornado.gen.coroutine
-    def handle_stream(stream):
-        yield stream.write('hello')
+    def handle_stream(self, stream, address):
+        yield stream.write(self.response)
         stream.close()
 
 
@@ -144,3 +148,42 @@ def test_check_failure(self):
         with mock.patch.object(checker, 'TIMEOUT', 1):
             response = yield checker.check_tcp("foo", self.unlistened_port, None, io_loop=self.io_loop, query_params="")
             self.assertEqual(response[0], 503)
+
+
+class TestRedisChecker(tornado.testing.AsyncTestCase):
+    def setUp(self):
+        super(TestRedisChecker, self).setUp()
+        socket, port = tornado.testing.bind_unused_port()
+        self.server = TestServer(io_loop=self.io_loop)
+        self.server.add_socket(socket)
+        self.socket = socket
+        self.port = port
+        unlistened_socket, unlistened_port = bind_synchronous_unused_port()
+        self.unlistened_socket = unlistened_socket
+        self.unlistened_port = unlistened_port
+
+    def tearDown(self):
+        super(TestRedisChecker, self).tearDown()
+        try:
+            self.server.stop()
+            self.socket.close()
+        except Exception:
+            pass
+
+    @tornado.testing.gen_test
+    def test_check_success(self):
+        with mock.patch.object(self.server, 'response', b'+PONG\r\n'):
+            response = yield checker.check_redis("foo", self.port, None, io_loop=self.io_loop, query_params="")
+            self.assertEqual(200, response[0], response[1])
+
+    @tornado.testing.gen_test
+    def test_check_error(self):
+        with mock.patch.object(self.server, 'response', b'WAT\r\n'):
+            response = yield checker.check_redis("foo", self.port, None, io_loop=self.io_loop, query_params="")
+            self.assertEqual(500, response[0])
+
+    @tornado.testing.gen_test
+    def test_check_timeout(self):
+        with mock.patch.object(checker, 'TIMEOUT', 1):
+            response = yield checker.check_tcp("foo", self.unlistened_port, None, io_loop=self.io_loop, query_params="")
+            self.assertEqual(response[0], 503)