fix(error): properly handle the 404 errors (#43)

Signed-off-by: Nicolas Crocfer <nicolas.crocfer@corp.ovh.com>
ovh · May 11, 2020 · 5e81a80 · 5e81a80
1 parent ba30a92
commit 5e81a80
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 1 deletion.
diff --git a/director/__init__.py b/director/__init__.py
@@ -4,7 +4,7 @@
 from functools import partial
 from pathlib import Path
 
-from flask import Flask, Blueprint, jsonify
+from flask import Flask, Blueprint, jsonify, request, render_template
 from flask_json_schema import JsonValidationError
 
 from director.api import api_bp
@@ -56,6 +56,9 @@ def create_app(
         )
     )
 
+    # Error handler
+    app.register_error_handler(404, lambda e: handle_not_found(e))
+
     # Init extensions
     db.init_app(app)
     db.app = app
@@ -86,6 +89,15 @@ def create_app(
     return app
 
 
+def handle_not_found(e):
+    """
+    Handle the "404 Not found" error in API and HTML.
+    """
+    if request.path.startswith("/api"):
+        return jsonify(error=str(e)), 404
+    return render_template("404.html", error=e), 404
+
+
 # Import director's submodules
 def import_submodules(package, modules_to_import):
     """

diff --git a/director/templates/404.html b/director/templates/404.html
@@ -0,0 +1,8 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <title>404 Not Found</title>
+</head>
+    <body>%% error %%</body>
+</html>
diff --git a/tests/test_views.py b/tests/test_views.py
@@ -0,0 +1,30 @@
+def test_view(client):
+    resp = client.get("/")
+    assert resp.status_code == 200
+    assert not resp.is_json
+    assert resp.mimetype == "text/html"
+    assert "<title>Celery Director</title>" in str(resp.data)
+
+
+def test_404_view(client):
+    resp = client.get("/notfound")
+    assert resp.status_code == 404
+    assert not resp.is_json
+    assert resp.mimetype == "text/html"
+    assert "<title>404 Not Found</title>" in str(resp.data)
+
+
+def test_api(client):
+    resp = client.get("/api/ping")
+    assert resp.status_code == 200
+    assert resp.is_json
+    assert resp.mimetype == "application/json"
+    assert resp.json == {"message": "pong"}
+
+
+def test_404_api(client):
+    resp = client.get("/api/notfound")
+    assert resp.status_code == 404
+    assert resp.is_json
+    assert resp.mimetype == "application/json"
+    assert "404 Not Found" in resp.json["error"]