Skip to content

Commit 717ec5b

Browse files
committed
Leverage centralized checks (#85)
1 parent c8e2b19 commit 717ec5b

File tree

4 files changed

+19
-79
lines changed

4 files changed

+19
-79
lines changed

src/dockerflow/fastapi/__init__.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
from fastapi import APIRouter
22
from fastapi.routing import APIRoute
33

4-
from .checks import register_heartbeat_check # noqa
54
from .views import heartbeat, lbheartbeat, version
65

76
router = APIRouter(

src/dockerflow/fastapi/checks.py

Lines changed: 0 additions & 52 deletions
This file was deleted.

src/dockerflow/fastapi/views.py

Lines changed: 12 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -5,35 +5,29 @@
55
from dockerflow import checks
66

77
from ..version import get_version
8-
from .checks import run_heartbeat_checks
98

109

1110
def lbheartbeat():
1211
return {"status": "ok"}
1312

1413

1514
def heartbeat(response: Response):
16-
check_results = run_heartbeat_checks()
17-
details = {}
18-
statuses = {}
19-
level = 0
20-
21-
for name, detail in check_results:
22-
statuses[name] = detail.status
23-
level = max(level, detail.level)
24-
if detail.level > 0:
25-
details[name] = detail
26-
27-
if level < checks.ERROR:
15+
check_results = checks.run_checks(
16+
checks.get_checks().items(),
17+
)
18+
19+
payload = {
20+
"status": checks.level_to_text(check_results.level),
21+
"checks": check_results.statuses,
22+
"details": check_results.details,
23+
}
24+
25+
if check_results.level < checks.ERROR:
2826
response.status_code = 200
2927
else:
3028
response.status_code = 500
3129

32-
return {
33-
"status": checks.level_to_text(level),
34-
"checks": statuses,
35-
"details": details,
36-
}
30+
return payload
3731

3832

3933
def version(request: Request):

tests/fastapi/test_fastapi.py

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,7 @@
88
from fastapi import FastAPI
99
from fastapi.testclient import TestClient
1010

11-
from dockerflow.checks import Error
12-
from dockerflow.fastapi import register_heartbeat_check
11+
from dockerflow import checks
1312
from dockerflow.fastapi import router as dockerflow_router
1413
from dockerflow.fastapi.middleware import MozlogRequestSummaryLogger
1514

@@ -96,9 +95,9 @@ def test_version_default(client, mocker):
9695

9796

9897
def test_heartbeat_get(client):
99-
@register_heartbeat_check
98+
@checks.register
10099
def return_error():
101-
return [Error("BOOM", id="foo")]
100+
return [checks.Error("BOOM", id="foo")]
102101

103102
response = client.get("/__heartbeat__")
104103
assert response.status_code == 500
@@ -116,18 +115,18 @@ def return_error():
116115

117116

118117
def test_heartbeat_head(client):
119-
@register_heartbeat_check
118+
@checks.register
120119
def return_error():
121-
return [Error("BOOM", id="foo")]
120+
return [checks.Error("BOOM", id="foo")]
122121

123122
response = client.head("/__heartbeat__")
124123
assert response.content == b""
125124

126125

127126
def test_heartbeat_custom_name(client):
128-
@register_heartbeat_check(name="my_check_name")
127+
@checks.register(name="my_check_name")
129128
def return_error():
130-
return [Error("BOOM", id="foo")]
129+
return [checks.Error("BOOM", id="foo")]
131130

132131
response = client.get("/__heartbeat__")
133132
assert response.json()["checks"]["my_check_name"]

0 commit comments

Comments
 (0)