Merge pull request #1721 from PrefectHQ/jawnsy/chore-use-status-const…

…ants

Use status constants instead of hardcoded values

src/prefect/cli/cloud.py

@@ -6,6 +6,7 @@
 import httpx
 import readchar
 import typer
+from fastapi import status
 from rich.live import Live
 from rich.table import Table
 
@@ -94,7 +95,10 @@ async def get(self, route, **kwargs):
             res = await self._client.get(route, **kwargs)
             res.raise_for_status()
         except httpx.HTTPStatusError as exc:
-            if exc.response.status_code in (401, 403):
+            if exc.response.status_code in (
+                status.HTTP_401_UNAUTHORIZED,
+                status.HTTP_403_FORBIDDEN,
+            ):
                 exit_with_error(
                     "Unable to authenticate. Please ensure your credentials are correct."
                 )

src/prefect/cli/flow_run.py

@@ -4,10 +4,9 @@
 from typing import List
 from uuid import UUID
 
-import fastapi
 import httpx
 import pendulum
-import typer
+from fastapi import status
 from rich.pretty import Pretty
 from rich.table import Table
 
@@ -32,7 +31,7 @@ async def inspect(id: UUID):
         try:
             flow_run = await client.read_flow_run(id)
         except httpx.HTTPStatusError as exc:
-            if exc.response.status_code == fastapi.status.HTTP_404_NOT_FOUND:
+            if exc.response.status_code == status.HTTP_404_NOT_FOUND:
                 exit_with_error(f"Flow run {id!r} not found!")
             else:
                 raise

src/prefect/client.py

@@ -38,7 +38,7 @@
 import httpx
 import pydantic
 from asgi_lifespan import LifespanManager
-from fastapi import FastAPI
+from fastapi import FastAPI, status
 
 import prefect
 import prefect.exceptions
@@ -594,7 +594,7 @@ async def read_concurrency_limit_by_tag(
                 f"/concurrency_limits/tag/{tag}",
             )
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 404:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
                 raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
             else:
                 raise
@@ -653,7 +653,7 @@ async def delete_concurrency_limit_by_tag(
                 f"/concurrency_limits/tag/{tag}",
             )
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 404:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
                 raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
             else:
                 raise
@@ -695,7 +695,7 @@ async def create_work_queue(
         try:
             response = await self._client.post("/work_queues/", json=data)
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 409:
+            if e.response.status_code == status.HTTP_409_CONFLICT:
                 raise prefect.exceptions.ObjectAlreadyExists(http_exc=e) from e
             else:
                 raise
@@ -742,7 +742,7 @@ async def update_work_queue(self, id: UUID, **kwargs):
         try:
             await self._client.patch(f"/work_queues/{id}", json=data)
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 404:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
                 raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
             else:
                 raise
@@ -779,7 +779,7 @@ async def get_runs_in_work_queue(
                 json=json_data,
             )
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 404:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
                 raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
             else:
                 raise
@@ -805,7 +805,7 @@ async def read_work_queue(
         try:
             response = await self._client.get(f"/work_queues/{id}")
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 404:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
                 raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
             else:
                 raise
@@ -853,7 +853,7 @@ async def delete_work_queue_by_id(
                 f"/work_queues/{id}",
             )
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 404:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
                 raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
             else:
                 raise
@@ -882,7 +882,7 @@ async def create_block(
                 json=payload,
             )
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 409:
+            if e.response.status_code == status.HTTP_409_CONFLICT:
                 raise prefect.exceptions.ObjectAlreadyExists(http_exc=e) from e
             else:
                 raise
@@ -897,7 +897,7 @@ async def read_block_spec_by_name(
         try:
             response = await self._client.get(f"block_specs/{name}/versions/{version}")
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 404:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
                 raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
             else:
                 raise
@@ -1076,7 +1076,7 @@ async def read_deployment_by_name(
         try:
             response = await self._client.get(f"/deployments/name/{name}")
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 404:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
                 raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
             else:
                 raise
@@ -1214,7 +1214,7 @@ async def set_default_storage_block(self, block_id: UUID):
         try:
             await self._client.post(f"/blocks/{block_id}/set_default_storage_block")
         except httpx.HTTPStatusError as e:
-            if e.response.status_code == 404:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
                 raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
             else:
                 raise

src/prefect/orion/api/block_specs.py

@@ -4,20 +4,9 @@
 from typing import List
 from uuid import UUID
 
-import pendulum
 import sqlalchemy as sa
-from fastapi import (
-    Body,
-    Depends,
-    HTTPException,
-    Path,
-    Query,
-    Response,
-    responses,
-    status,
-)
-
-from prefect import settings
+from fastapi import Body, Depends, HTTPException, Path, Response, status
+
 from prefect.orion import models, schemas
 from prefect.orion.api import dependencies
 from prefect.orion.database.dependencies import provide_database_interface

src/prefect/orion/api/blocks.py

@@ -4,11 +4,9 @@
 from typing import List, Optional
 from uuid import UUID
 
-import pendulum
 import sqlalchemy as sa
 from fastapi import Body, Depends, HTTPException, Path, Response, responses, status
 
-from prefect import settings
 from prefect.orion import models, schemas
 from prefect.orion.api import dependencies
 from prefect.orion.database.dependencies import provide_database_interface
@@ -64,7 +62,8 @@ async def read_block_by_id(
     model = await models.blocks.read_block_by_id(session=session, block_id=block_id)
     if not model:
         return responses.JSONResponse(
-            status_code=404, content={"message": "Block not found"}
+            status_code=status.HTTP_404_NOT_FOUND,
+            content={"message": "Block not found"},
         )
     return await schemas.core.Block.from_orm_model(session=session, orm_block=model)
 

src/prefect/orion/api/concurrency_limits.py

@@ -11,7 +11,6 @@
 import prefect.orion.api.dependencies as dependencies
 import prefect.orion.models as models
 import prefect.orion.schemas as schemas
-import prefect.settings
 from prefect.orion.database.dependencies import provide_database_interface
 from prefect.orion.database.interface import OrionDBInterface
 from prefect.orion.utilities.server import OrionRouter

src/prefect/orion/api/logs.py

@@ -5,13 +5,11 @@
 from typing import List
 
 import sqlalchemy as sa
-from fastapi import Body, Depends, Response
-from starlette import status
+from fastapi import Body, Depends, Response, status
 
 import prefect.orion.api.dependencies as dependencies
 import prefect.orion.models as models
 import prefect.orion.schemas as schemas
-import prefect.settings
 from prefect.orion.utilities.server import OrionRouter
 
 router = OrionRouter(prefix="/logs", tags=["Logs"])

src/prefect/orion/api/server.py

@@ -46,7 +46,7 @@ class SPAStaticFiles(StaticFiles):
     async def get_response(self, path: str, scope):
         response = await super().get_response(path, scope)
 
-        if response.status_code == 404:
+        if response.status_code == status.HTTP_404_NOT_FOUND:
             response = await super().get_response("./index.html", scope)
 
         return response

src/prefect/utilities/httpx.py

@@ -1,11 +1,8 @@
-import typing
-
 import anyio
 import httpx
+from fastapi import status
 from httpx._models import Response
 
-from prefect.settings import PREFECT_API_REQUEST_TIMEOUT
-
 
 class PrefectHttpxClient(httpx.AsyncClient):
     RETRY_MAX = 5
@@ -20,7 +17,10 @@ async def send(self, *args, **kwargs) -> Response:
         retry_count = 0
         response = await self._httpx_send(*args, **kwargs)
 
-        while response.status_code == 429 and retry_count < self.RETRY_MAX:
+        while (
+            response.status_code == status.HTTP_429_TOO_MANY_REQUESTS
+            and retry_count < self.RETRY_MAX
+        ):
             retry_count += 1
             # respect CloudFlare conventions for handling rate-limit response headers
             # see https://support.cloudflare.com/hc/en-us/articles/115001635128-Configuring-Rate-Limiting-from-UI

tests/cli/test_config.py

@@ -7,7 +7,7 @@
 import prefect.settings
 from prefect.cli import app
 from prefect.context import profile
-from prefect.utilities.testing import temporary_settings
+from prefect.testing.utilities import temporary_settings
 
 """
 Testing Typer tutorial here: https://typer.tiangolo.com/tutorial/testing/

tests/fixtures/storage.py

@@ -3,19 +3,17 @@
 """
 import subprocess
 import sys
-import time
 from typing import Any, Optional
 
 import anyio
 import httpx
 import pytest
-from fastapi import Body, FastAPI
+from fastapi import Body, FastAPI, status
 from fastapi.exceptions import RequestValidationError
 
 from prefect.blocks.storage import KVServerStorageBlock, LocalStorageBlock
 from prefect.orion import models
 from prefect.orion.api.server import validation_exception_handler
-from prefect.orion.schemas.actions import BlockCreate, BlockSpecCreate
 
 
 @pytest.fixture
@@ -89,7 +87,7 @@ async def run_storage_server():
                     except httpx.ConnectError:
                         pass
                     else:
-                        if response.status_code == 200:
+                        if response.status_code == status.HTTP_200_OK:
                             break
                     await anyio.sleep(0.1)
             if response:

tests/orion/api/test_admin.py

@@ -1,24 +1,26 @@
+from fastapi import status
+
 import prefect
 from prefect.orion import models
 
 
 async def test_hello_world(client):
     response = await client.get("/admin/hello")
-    assert response.status_code == 200
+    assert response.status_code == status.HTTP_200_OK
     assert response.json() == "👋"
 
 
 async def test_version(client):
     response = await client.get("/admin/version")
-    assert response.status_code == 200
+    assert response.status_code == status.HTTP_200_OK
     assert prefect.__version__
     assert response.json() == prefect.__version__
 
 
 class TestSettings:
     async def test_read_settings(self, client):
         response = await client.get("/admin/settings")
-        assert response.status_code == 200
+        assert response.status_code == status.HTTP_200_OK
         parsed_settings = prefect.settings.Settings.parse_obj(response.json())
         prefect_settings = prefect.settings.get_current_settings()
 
@@ -38,7 +40,7 @@ async def test_clear_database(
             assert await count(session) > 0
 
         response = await client.post("/admin/database/clear", json=dict(confirm=True))
-        assert response.status_code == 204
+        assert response.status_code == status.HTTP_204_NO_CONTENT
 
         for count in [
             models.flows.count_flows,
@@ -50,21 +52,21 @@ async def test_clear_database(
 
     async def test_clear_database_requires_confirmation(self, client):
         response = await client.post("/admin/database/clear")
-        assert response.status_code == 400
+        assert response.status_code == status.HTTP_400_BAD_REQUEST
 
         response = await client.post("/admin/database/clear", json=dict(confirm=False))
-        assert response.status_code == 400
+        assert response.status_code == status.HTTP_400_BAD_REQUEST
 
     async def test_drop_database_requires_confirmation(self, client):
         response = await client.post("/admin/database/drop")
-        assert response.status_code == 400
+        assert response.status_code == status.HTTP_400_BAD_REQUEST
 
         response = await client.post("/admin/database/drop", json=dict(confirm=False))
-        assert response.status_code == 400
+        assert response.status_code == status.HTTP_400_BAD_REQUEST
 
     async def test_create_database_requires_confirmation(self, client):
         response = await client.post("/admin/database/create")
-        assert response.status_code == 400
+        assert response.status_code == status.HTTP_400_BAD_REQUEST
 
         response = await client.post("/admin/database/create", json=dict(confirm=False))
-        assert response.status_code == 400
+        assert response.status_code == status.HTTP_400_BAD_REQUEST