Skip to content

Commit

Permalink
Allowing for async checks in fastApi. Cleaning up old code. (#114)
Browse files Browse the repository at this point in the history
* Allowing for async checks in fastApi. Cleaning up old code. Fixes #106
* Added a few unit tests for fast_api
  • Loading branch information
alexcottner authored Apr 24, 2024
1 parent 0424734 commit 391f9ba
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 17 deletions.
17 changes: 2 additions & 15 deletions src/dockerflow/checks/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,19 +17,6 @@
_REGISTERED_CHECKS = {}


def _iscoroutinefunction_or_partial(obj):
"""
Determine if the provided object is a coroutine function or a partial function
that wraps a coroutine function.
This function should be removed when we drop support for Python 3.7, as this is
handled directly by `inspect.iscoroutinefunction` in Python 3.8.
"""
while isinstance(obj, functools.partial):
obj = obj.func
return inspect.iscoroutinefunction(obj)


def register(func=None, name=None):
"""
Register a check callback to be executed from
Expand All @@ -43,7 +30,7 @@ def register(func=None, name=None):

logger.debug("Register Dockerflow check %s", name)

if _iscoroutinefunction_or_partial(func):
if inspect.iscoroutinefunction(func):

@functools.wraps(func)
async def decorated_function_asyc(*args, **kwargs):
Expand Down Expand Up @@ -116,7 +103,7 @@ class ChecksResults:

async def _run_check_async(check):
name, check_fn = check
if _iscoroutinefunction_or_partial(check_fn):
if inspect.iscoroutinefunction(check_fn):
errors = await check_fn()
else:
loop = asyncio.get_event_loop()
Expand Down
4 changes: 2 additions & 2 deletions src/dockerflow/fastapi/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,12 @@ def lbheartbeat():
return {"status": "ok"}


def heartbeat(request: Request, response: Response):
async def heartbeat(request: Request, response: Response):
FAILED_STATUS_CODE = int(
getattr(request.app.state, "DOCKERFLOW_HEARTBEAT_FAILED_STATUS_CODE", "500")
)

check_results = checks.run_checks(
check_results = await checks.run_checks_async(
checks.get_checks().items(),
)

Expand Down
47 changes: 47 additions & 0 deletions tests/fastapi/test_fastapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,53 @@ def return_error():
},
}

def test_heartbeat_sync(client):
@checks.register
def sync_ok():
return []

response = client.get("/__heartbeat__")
assert response.status_code == 200
assert response.json() == {
"status": "ok",
"checks": {"sync_ok": "ok"},
"details": {},
}


def test_heartbeat_async(client):
@checks.register
async def async_ok():
return []

response = client.get("/__heartbeat__")
assert response.status_code == 200
assert response.json() == {
"status": "ok",
"checks": {"async_ok": "ok"},
"details": {},
}


def test_heartbeat_mixed_sync(client):
@checks.register
def sync_ok():
return []
@checks.register
async def async_ok():
return []

response = client.get("/__heartbeat__")
assert response.status_code == 200
assert response.json() == {
"status": "ok",
"checks": {
"sync_ok": "ok",
"async_ok": "ok",
},
"details": {},
}


def test_heartbeat_head(client):
@checks.register
Expand Down

0 comments on commit 391f9ba

Please sign in to comment.