feature: avoid having magic number for status codes

Signed-off-by: F.N. Claessen <[email protected]>
FlexMeasures · Jun 5, 2024 · c7e58cb · c7e58cb
1 parent 05064ae
commit c7e58cb
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 7 deletions.
diff --git a/flexmeasures/api/v3_0/assets.py b/flexmeasures/api/v3_0/assets.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from datetime import datetime, timedelta
+from http import HTTPStatus
 import json
 
 from flask import current_app
@@ -165,7 +166,12 @@ def post(self, asset_data: dict):
 
     @route("/<id>", methods=["GET"])
     @use_kwargs(
-        {"asset": AssetIdField(data_key="id", status_if_not_found=404)}, location="path"
+        {
+            "asset": AssetIdField(
+                data_key="id", status_if_not_found=HTTPStatus.NOT_FOUND
+            )
+        },
+        location="path",
     )
     @permission_required_for_context("read", ctx_arg_name="asset")
     @as_json
@@ -203,7 +209,11 @@ def fetch_one(self, id, asset):
     @route("/<id>", methods=["PATCH"])
     @use_args(partial_asset_schema)
     @use_kwargs(
-        {"db_asset": AssetIdField(data_key="id", status_if_not_found=404)},
+        {
+            "db_asset": AssetIdField(
+                data_key="id", status_if_not_found=HTTPStatus.NOT_FOUND
+            )
+        },
         location="path",
     )
     @permission_required_for_context("update", ctx_arg_name="db_asset")
@@ -265,7 +275,12 @@ def patch(self, asset_data: dict, id: int, db_asset: GenericAsset):
 
     @route("/<id>", methods=["DELETE"])
     @use_kwargs(
-        {"asset": AssetIdField(data_key="id", status_if_not_found=404)}, location="path"
+        {
+            "asset": AssetIdField(
+                data_key="id", status_if_not_found=HTTPStatus.NOT_FOUND
+            )
+        },
+        location="path",
     )
     @permission_required_for_context("delete", ctx_arg_name="asset")
     @as_json
@@ -293,7 +308,11 @@ def delete(self, id: int, asset: GenericAsset):
 
     @route("/<id>/chart", strict_slashes=False)  # strict on next version? see #1014
     @use_kwargs(
-        {"asset": AssetIdField(data_key="id", status_if_not_found=404)},
+        {
+            "asset": AssetIdField(
+                data_key="id", status_if_not_found=HTTPStatus.NOT_FOUND
+            )
+        },
         location="path",
     )
     @use_kwargs(
@@ -323,7 +342,11 @@ def get_chart(self, id: int, asset: GenericAsset, **kwargs):
         "/<id>/chart_data", strict_slashes=False
     )  # strict on next version? see #1014
     @use_kwargs(
-        {"asset": AssetIdField(data_key="id", status_if_not_found=404)},
+        {
+            "asset": AssetIdField(
+                data_key="id", status_if_not_found=HTTPStatus.NOT_FOUND
+            )
+        },
         location="path",
     )
     @use_kwargs(

diff --git a/flexmeasures/data/schemas/generic_assets.py b/flexmeasures/data/schemas/generic_assets.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+from http import HTTPStatus
 
 from flask import abort
 from marshmallow import validates, ValidationError, fields, validates_schema
@@ -151,7 +152,7 @@ class Meta:
 class GenericAssetIdField(MarshmallowClickMixin, fields.Int):
     """Field that deserializes to a GenericAsset and serializes back to an integer."""
 
-    def __init__(self, status_if_not_found: int | None = None, *args, **kwargs):
+    def __init__(self, status_if_not_found: HTTPStatus | None = None, *args, **kwargs):
         self.status_if_not_found = status_if_not_found
         super().__init__(*args, **kwargs)
 
@@ -162,7 +163,7 @@ def _deserialize(self, value: int, attr, obj, **kwargs) -> GenericAsset:
         ).scalar_one_or_none()
         if generic_asset is None:
             message = f"No asset found with ID {value}."
-            if self.status_if_not_found == 404:
+            if self.status_if_not_found == HTTPStatus.NOT_FOUND:
                 raise abort(404, message)
             else:
                 raise FMValidationError(message)