Additional type checks, type hints and field descriptions

Author: Nadine-H

src/dstack/plugins/builtin/rest_plugin/_models.py

@@ -1,6 +1,7 @@
-from typing import Generic, TypeVar
+from typing import Generic, Optional, TypeVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated
 
 from dstack._internal.core.models.fleets import FleetSpec
 from dstack._internal.core.models.gateways import GatewaySpec
@@ -11,9 +12,9 @@
 
 
 class SpecApplyRequest(BaseModel, Generic[SpecType]):
-    user: str
-    project: str
-    spec: SpecType
+    user: Annotated[str, Field(description="The name of the user making the apply request")]
+    project: Annotated[str, Field(description="The name of the project the request is for")]
+    spec: Annotated[SpecType, Field(description="The spec to be applied")]
 
     # Override dict() to remove __orig_class__ attribute and avoid "TypeError: Object of type _GenericAlias is not JSON serializable"
     # error. This issue doesn't happen though when running the code in pytest, only when running the server.
@@ -30,8 +31,15 @@ def dict(self, *args, **kwargs):
 
 
 class SpecApplyResponse(BaseModel, Generic[SpecType]):
-    spec: SpecType
-    error: str | None = None
+    spec: Annotated[
+        SpecType,
+        Field(
+            description="The spec to apply, original spec if error otherwise original or mutated by plugin service if approved"
+        ),
+    ]
+    error: Annotated[
+        Optional[str], Field(description="Error message if request is rejected", min_length=1)
+    ] = None
 
 
 RunSpecResponse = SpecApplyResponse[RunSpec]

src/dstack/plugins/builtin/rest_plugin/_plugin.py

@@ -1,5 +1,6 @@
 import json
 import os
+from typing import Type
 
 import requests
 from pydantic import ValidationError
@@ -25,7 +26,7 @@
 logger = get_plugin_logger(__name__)
 
 PLUGIN_SERVICE_URI_ENV_VAR_NAME = "DSTACK_PLUGIN_SERVICE_URI"
-PLUGIN_REQUEST_TIMEOUT = 8  # in seconds
+PLUGIN_REQUEST_TIMEOUT_SEC = 8
 
 
 class CustomApplyPolicy(ApplyPolicy):
@@ -50,7 +51,7 @@ def _call_plugin_service(self, spec_request: SpecApplyRequest, endpoint: str) ->
                 f"{self._plugin_service_uri}{endpoint}",
                 json=spec_request.dict(),
                 headers={"accept": "application/json", "Content-Type": "application/json"},
-                timeout=PLUGIN_REQUEST_TIMEOUT,
+                timeout=PLUGIN_REQUEST_TIMEOUT_SEC,
             )
             response.raise_for_status()
             spec_json = json.loads(response.text)
@@ -66,7 +67,15 @@ def _call_plugin_service(self, spec_request: SpecApplyRequest, endpoint: str) ->
             logger.error("Request to the plugin service failed: %s", e)
             raise ServerClientError("Request to the plugin service failed")
 
-    def _on_apply(self, request_cls, response_cls, endpoint, user, project, spec):
+    def _on_apply(
+        self,
+        request_cls: Type[SpecApplyRequest],
+        response_cls: Type[SpecApplyResponse],
+        endpoint: str,
+        user: str,
+        project: str,
+        spec: ApplySpec,
+    ) -> ApplySpec:
         try:
             spec_request = request_cls(user=user, project=project, spec=spec)
             spec_json = self._call_plugin_service(spec_request, endpoint)

src/tests/plugins/test_rest_plugin.py

@@ -1,5 +1,6 @@
 import json
 import os
+from contextlib import nullcontext as does_not_raise
 from unittest import mock
 from unittest.mock import Mock
 
@@ -165,3 +166,42 @@ async def test_on_apply_plugin_service_returns_invalid_spec(
         with mock.patch("requests.post", return_value=mock_response):
             with pytest.raises(ServerClientError):
                 policy.on_apply(user.name, project=project.name, spec=spec)
+
+    @pytest.mark.asyncio
+    @mock.patch.dict(os.environ, {PLUGIN_SERVICE_URI_ENV_VAR_NAME: "http://mock"})
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.parametrize(
+        "spec", ["run_spec", "fleet_spec", "volume_spec", "gateway_spec"], indirect=True
+    )
+    @pytest.mark.parametrize(
+        ("error", "expectation"),
+        [
+            pytest.param(None, does_not_raise(), id="error_none"),
+            pytest.param(
+                "",
+                pytest.raises(
+                    ServerClientError, match="Plugin service returned an invalid response"
+                ),
+                id="error_empty_str",
+            ),
+            pytest.param(
+                "validation failed",
+                pytest.raises(
+                    ServerClientError, match="Apply request rejected: validation failed"
+                ),
+                id="error_non_empty_str",
+            ),
+        ],
+    )
+    async def test_on_apply_plugin_service_error_handling(
+        self, test_db, user, project, spec, error, expectation
+    ):
+        policy = CustomApplyPolicy()
+        mock_response = Mock()
+        response_dict = {"spec": spec.dict(), "error": error}
+        mock_response.text = json.dumps(response_dict)
+        mock_response.raise_for_status = Mock()
+        with mock.patch("requests.post", return_value=mock_response):
+            with expectation:
+                result = policy.on_apply(user=user.name, project=project.name, spec=spec)
+                assert result == type(spec)(**response_dict["spec"])