diff --git a/adsmutils/__init__.py b/adsmutils/__init__.py index 35d4d48..6b3b5bc 100644 --- a/adsmutils/__init__.py +++ b/adsmutils/__init__.py @@ -24,6 +24,7 @@ import inspect from cloghandler import ConcurrentRotatingFileHandler from flask import Flask +from flask import Response from pythonjsonlogger import jsonlogger from logging import Formatter import flask @@ -293,6 +294,9 @@ def __init__(self, app_name, *args, **kwargs): self.client.mount('http://', http_adapter) self.before_request_funcs.setdefault(None, []).append(self._before_request) + self.add_url_rule('/ready', 'ready', self.ready) + self.add_url_rule('/alive', 'alive', self.alive) + def _before_request(self): if flask.has_request_context(): # New request will contain also key information from the original request @@ -318,6 +322,28 @@ def close_app(self): self.db = None self.logger = None + def ready(self, key='ready'): + """Endpoint /ready to signal that the application is ready to receive requests""" + if self._db_failure(): + return Response(json.dumps({key: False}), mimetype='application/json', status=503) + else: + return Response(json.dumps({key: True}), mimetype='application/json', status=200) + + def alive(self): + """Endpoint /alive to signal that the application is healthy""" + return self.ready(key="alive") + + def _db_failure(self): + if self.db is None: + return False + else: + with self.session_scope() as session: + try: + session.execute('SELECT 1') + return False + except: + return True + @contextmanager def session_scope(self): diff --git a/adsmutils/tests/base.py b/adsmutils/tests/base.py new file mode 100644 index 0000000..e9b778a --- /dev/null +++ b/adsmutils/tests/base.py @@ -0,0 +1,58 @@ +# -*- coding: utf-8 -*- + +import testing.postgresql +from flask_testing import TestCase +from adsmutils import ADSFlask + +class TestCaseDatabase(TestCase): + + postgresql_url_dict = { + 'port': 1234, + 'host': '127.0.0.1', + 'user': 'postgres', + 'database': 'test' + } + postgresql_url = 'postgresql://{user}@{host}:{port}/{database}' \ + .format( + user=postgresql_url_dict['user'], + host=postgresql_url_dict['host'], + port=postgresql_url_dict['port'], + database=postgresql_url_dict['database'] + ) + + def create_app(self): + '''Start the wsgi application''' + local_config = { + 'SQLALCHEMY_DATABASE_URI': self.postgresql_url, + 'SQLALCHEMY_ECHO': False, + 'TESTING': True, + 'PROPAGATE_EXCEPTIONS': True, + 'TRAP_BAD_REQUEST_ERRORS': True + } + app = ADSFlask(__name__, static_folder=None, local_config=local_config) + return app + + @classmethod + def setUpClass(cls): + cls.postgresql = \ + testing.postgresql.Postgresql(**cls.postgresql_url_dict) + + @classmethod + def tearDownClass(cls): + cls.postgresql.stop() + + def setUp(self): + pass + + def tearDown(self): + self.app.db.session.remove() + self.app.db.drop_all() + +class TestCase(TestCase): + + def create_app(self): + '''Start the wsgi application''' + local_config = { } + app = ADSFlask(__name__, static_folder=None, local_config=local_config) + return app + diff --git a/adsmutils/tests/test_service.py b/adsmutils/tests/test_service.py new file mode 100644 index 0000000..e226ea8 --- /dev/null +++ b/adsmutils/tests/test_service.py @@ -0,0 +1,53 @@ +# -*- coding: utf-8 -*- + +from base import TestCase, TestCaseDatabase +import mock + +class TestServices(TestCase): + + def test_readiness_probe(self): + '''Tests for the existence of a /ready route, and that it returns properly + formatted JSON data''' + r = self.client.get('/ready') + self.assertEqual(r.status_code, 200) + self.assertEqual(r.json['ready'], True) + + def test_liveliness_probe(self): + '''Tests for the existence of a /alive route, and that it returns properly + formatted JSON data''' + r = self.client.get('/alive') + self.assertEqual(r.status_code, 200) + self.assertEqual(r.json['alive'], True) + + +class TestServicesWithDatabase(TestCaseDatabase): + + def test_readiness_probe(self): + '''Tests for the existence of a /ready route, and that it returns properly + formatted JSON data''' + r = self.client.get('/ready') + self.assertEqual(r.status_code, 200) + self.assertEqual(r.json['ready'], True) + + def test_liveliness_probe(self): + '''Tests for the existence of a /alive route, and that it returns properly + formatted JSON data''' + r = self.client.get('/alive') + self.assertEqual(r.status_code, 200) + self.assertEqual(r.json['alive'], True) + + def test_readiness_probe_with_db_failure(self): + '''Tests for the existence of a /ready route, and that it returns properly + formatted JSON data when database connection has been lost''' + self.app._db_failure = mock.MagicMock(return_value=True) + r = self.client.get('/ready') + self.assertEqual(r.status_code, 503) + self.assertEqual(r.json['ready'], False) + + def test_liveliness_probe_with_db_failure(self): + '''Tests for the existence of a /alive route, and that it returns properly + formatted JSON data when database connection has been lost''' + self.app._db_failure = mock.MagicMock(return_value=True) + r = self.client.get('/alive') + self.assertEqual(r.status_code, 503) + self.assertEqual(r.json['alive'], False) diff --git a/dev-requirements.txt b/dev-requirements.txt index 575e920..fdbc9ed 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -1,6 +1,9 @@ pytest==2.8.2 +Flask-Testing==0.7.1 coveralls==1.1.0 httpretty==0.8.10 mock==1.3.0 coverage==4.0.1 pytest-cov==2.2.0 +testing.postgresql==1.2.1 +psycopg2==2.7.5