changes based on PR comments

package/kedro_viz/api/apps.py

@@ -14,7 +14,7 @@
 from jinja2 import Environment, FileSystemLoader
 
 from kedro_viz import __version__
-from kedro_viz.api.rest.responses.common import EnhancedORJSONResponse
+from kedro_viz.api.rest.responses.utils import EnhancedORJSONResponse
 from kedro_viz.integrations.kedro import telemetry as kedro_telemetry
 
 from .graphql.router import router as graphql_router

package/kedro_viz/api/rest/responses/base.py

@@ -0,0 +1,18 @@
+"""`kedro_viz.api.rest.responses.base` contains base
+response classes and utility functions for the REST endpoints"""
+
+# pylint: disable=missing-class-docstring
+
+import abc
+import logging
+from pydantic import BaseModel, ConfigDict
+
+logger = logging.getLogger(__name__)
+
+
+class APINotFoundResponse(BaseModel):
+    message: str
+
+
+class BaseAPIResponse(BaseModel, abc.ABC):
+    model_config = ConfigDict(from_attributes=True)

package/kedro_viz/api/rest/responses/metadata.py

@@ -7,7 +7,7 @@
 
 from pydantic import ConfigDict
 
-from kedro_viz.api.rest.responses.common import BaseAPIResponse
+from kedro_viz.api.rest.responses.base import BaseAPIResponse
 from kedro_viz.api.rest.utils import get_package_compatibilities
 from kedro_viz.models.metadata import Metadata, PackageCompatibility
 

package/kedro_viz/api/rest/responses/nodes.py

@@ -4,15 +4,12 @@
 # pylint: disable=missing-class-docstring,invalid-name
 
 import logging
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from fastapi.responses import JSONResponse
 from pydantic import ConfigDict
 
-from kedro_viz.api.rest.responses.common import (
-    BaseAPIResponse,
-    write_api_response_to_fs,
-)
+from kedro_viz.api.rest.responses.base import BaseAPIResponse
 from kedro_viz.data_access import data_access_manager
 from kedro_viz.models.flowchart import (
     DataNode,
@@ -125,22 +122,3 @@ def get_node_metadata_response(node_id: str):
         return TranscodedDataNodeMetadata(transcoded_data_node=node)
 
     return ParametersNodeMetadata(parameters_node=node)
-
-
-def save_api_node_response_to_fs(
-    nodes_path: str, remote_fs: Any, is_all_previews_enabled: bool
-):
-    """Saves API /nodes/{node} response to a directory."""
-    # Set if preview is enabled/disabled for all data nodes
-    DataNodeMetadata.set_is_all_previews_enabled(is_all_previews_enabled)
-
-    for nodeId in data_access_manager.nodes.get_node_ids():
-        try:
-            write_api_response_to_fs(
-                f"{nodes_path}/{nodeId}", get_node_metadata_response(nodeId), remote_fs
-            )
-        except Exception as exc:  # pragma: no cover
-            logger.exception(
-                "Failed to save node data for node ID %s. Error: %s", nodeId, str(exc)
-            )
-            raise exc

package/kedro_viz/api/rest/responses/pipelines.py

@@ -5,16 +5,13 @@
 
 import json
 import logging
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from fastapi.responses import JSONResponse
 from pydantic import ConfigDict
 
-from kedro_viz.api.rest.responses.common import (
-    BaseAPIResponse,
-    get_encoded_response,
-    write_api_response_to_fs,
-)
+from kedro_viz.api.rest.responses.base import BaseAPIResponse
+from kedro_viz.api.rest.responses.utils import get_encoded_response
 from kedro_viz.data_access import data_access_manager
 
 logger = logging.getLogger(__name__)
@@ -193,33 +190,6 @@ def get_pipeline_response(
     )
 
 
-def save_api_main_response_to_fs(main_path: str, remote_fs: Any):
-    """Saves API /main response to a directory."""
-    try:
-        write_api_response_to_fs(main_path, get_pipeline_response(), remote_fs)
-    except Exception as exc:  # pragma: no cover
-        logger.exception("Failed to save default response. Error: %s", str(exc))
-        raise exc
-
-
-def save_api_pipeline_response_to_fs(pipelines_path: str, remote_fs: Any):
-    """Saves API /pipelines/{pipeline} response to a directory."""
-    for pipelineId in data_access_manager.registered_pipelines.get_pipeline_ids():
-        try:
-            write_api_response_to_fs(
-                f"{pipelines_path}/{pipelineId}",
-                get_pipeline_response(pipelineId),
-                remote_fs,
-            )
-        except Exception as exc:  # pragma: no cover
-            logger.exception(
-                "Failed to save pipeline data for pipeline ID %s. Error: %s",
-                pipelineId,
-                str(exc),
-            )
-            raise exc
-
-
 def get_kedro_project_json_data():
     """Decodes the default response and returns the Kedro project JSON data.
     This will be used in VSCode extension to get current Kedro project data."""

package/kedro_viz/api/rest/responses/save_responses.py

@@ -0,0 +1,97 @@
+"""`kedro_viz.api.rest.responses.deploy` contains response classes
+and utility functions for the `/deploy` REST endpoint"""
+
+# pylint: disable=invalid-name
+
+import logging
+from typing import Any
+
+from kedro_viz.api.rest.responses.nodes import get_node_metadata_response
+from kedro_viz.api.rest.responses.pipelines import get_pipeline_response
+from kedro_viz.api.rest.responses.utils import get_encoded_response
+from kedro_viz.data_access import data_access_manager
+from kedro_viz.models.flowchart import DataNodeMetadata
+
+logger = logging.getLogger(__name__)
+
+
+def save_api_responses_to_fs(path: str, remote_fs: Any, is_all_previews_enabled: bool):
+    """Saves all Kedro Viz API responses to a directory."""
+    try:
+        logger.debug(
+            """Saving/Uploading api files to %s""",
+            path,
+        )
+
+        main_path = f"{path}/api/main"
+        nodes_path = f"{path}/api/nodes"
+        pipelines_path = f"{path}/api/pipelines"
+
+        if "file" in remote_fs.protocol:
+            remote_fs.makedirs(path, exist_ok=True)
+            remote_fs.makedirs(nodes_path, exist_ok=True)
+            remote_fs.makedirs(pipelines_path, exist_ok=True)
+
+        save_api_main_response_to_fs(main_path, remote_fs)
+        save_api_node_response_to_fs(nodes_path, remote_fs, is_all_previews_enabled)
+        save_api_pipeline_response_to_fs(pipelines_path, remote_fs)
+
+    except Exception as exc:  # pragma: no cover
+        logger.exception(
+            "An error occurred while preparing data for saving. Error: %s", str(exc)
+        )
+        raise exc
+
+
+def save_api_main_response_to_fs(main_path: str, remote_fs: Any):
+    """Saves API /main response to a directory."""
+    try:
+        write_api_response_to_fs(main_path, get_pipeline_response(), remote_fs)
+    except Exception as exc:  # pragma: no cover
+        logger.exception("Failed to save default response. Error: %s", str(exc))
+        raise exc
+
+
+def save_api_pipeline_response_to_fs(pipelines_path: str, remote_fs: Any):
+    """Saves API /pipelines/{pipeline} response to a directory."""
+    for pipeline_id in data_access_manager.registered_pipelines.get_pipeline_ids():
+        try:
+            write_api_response_to_fs(
+                f"{pipelines_path}/{pipeline_id}",
+                get_pipeline_response(pipeline_id),
+                remote_fs,
+            )
+        except Exception as exc:  # pragma: no cover
+            logger.exception(
+                "Failed to save pipeline data for pipeline ID %s. Error: %s",
+                pipeline_id,
+                str(exc),
+            )
+            raise exc
+
+
+def save_api_node_response_to_fs(
+    nodes_path: str, remote_fs: Any, is_all_previews_enabled: bool
+):
+    """Saves API /nodes/{node} response to a directory."""
+    # Set if preview is enabled/disabled for all data nodes
+    DataNodeMetadata.set_is_all_previews_enabled(is_all_previews_enabled)
+
+    for node_id in data_access_manager.nodes.get_node_ids():
+        try:
+            write_api_response_to_fs(
+                f"{nodes_path}/{node_id}", get_node_metadata_response(node_id), remote_fs
+            )
+        except Exception as exc:  # pragma: no cover
+            logger.exception(
+                "Failed to save node data for node ID %s. Error: %s", node_id, str(exc)
+            )
+            raise exc
+
+
+def write_api_response_to_fs(file_path: str, response: Any, remote_fs: Any):
+    """Get encoded responses and writes it to a file"""
+    encoded_response = get_encoded_response(response)
+
+    with remote_fs.open(file_path, "wb") as file:
+        file.write(encoded_response)

package/kedro_viz/api/rest/responses/common.py → package/kedro_viz/api/rest/responses/utils.py

@@ -1,28 +1,18 @@
-"""`kedro_viz.api.rest.responses.common` contains common
-response classes and utility functions for the REST endpoints"""
+"""`kedro_viz.api.rest.responses.utils` contains utility
+response classes and functions for the REST endpoints"""
 
 # pylint: disable=missing-class-docstring
 
-import abc
 import logging
 from typing import Any
 
 import orjson
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import ORJSONResponse
-from pydantic import BaseModel, ConfigDict
 
 logger = logging.getLogger(__name__)
 
 
-class APINotFoundResponse(BaseModel):
-    message: str
-
-
-class BaseAPIResponse(BaseModel, abc.ABC):
-    model_config = ConfigDict(from_attributes=True)
-
-
 class EnhancedORJSONResponse(ORJSONResponse):
     @staticmethod
     def encode_to_human_readable(content: Any) -> bytes:
@@ -49,11 +39,3 @@ def get_encoded_response(response: Any) -> bytes:
     )
 
     return encoded_response
-
-
-def write_api_response_to_fs(file_path: str, response: Any, remote_fs: Any):
-    """Get encoded responses and writes it to a file"""
-    encoded_response = get_encoded_response(response)
-
-    with remote_fs.open(file_path, "wb") as file:
-        file.write(encoded_response)

package/kedro_viz/api/rest/router.py

@@ -7,7 +7,7 @@
 from fastapi.responses import JSONResponse
 
 from kedro_viz.api.rest.requests import DeployerConfiguration
-from kedro_viz.api.rest.responses.common import APINotFoundResponse
+from kedro_viz.api.rest.responses.base import APINotFoundResponse
 from kedro_viz.api.rest.responses.metadata import (
     MetadataAPIResponse,
     get_metadata_response,

package/kedro_viz/integrations/deployment/base_deployer.py

@@ -12,7 +12,7 @@
 from packaging.version import parse
 
 from kedro_viz import __version__
-from kedro_viz.api.rest.responses.deploy import save_api_responses_to_fs
+from kedro_viz.api.rest.responses.save_responses import save_api_responses_to_fs
 from kedro_viz.integrations.kedro import telemetry as kedro_telemetry
 
 _HTML_DIR = Path(__file__).parent.parent.parent.absolute() / "html"

package/kedro_viz/server.py

@@ -126,7 +126,9 @@ def run_server(
         # we need to shift this feature outside of kedro viz run]
         if save_file:
             # pylint: disable=import-outside-toplevel
-            from kedro_viz.api.rest.responses.deploy import save_api_responses_to_fs
+            from kedro_viz.api.rest.responses.save_responses import (
+                save_api_responses_to_fs,
+            )
 
             save_api_responses_to_fs(save_file, fsspec.filesystem("file"), True)
 

package/tests/test_api/test_rest/test_responses/test_base.py

@@ -0,0 +1,10 @@
+from kedro_viz.api.rest.responses.base import APINotFoundResponse
+
+
+def test_api_not_found_response_valid_message():
+    response = APINotFoundResponse(message="Resource not found")
+    assert response.message == "Resource not found"
+
+    # Test that the model is serializable to a dictionary
+    serialized_response = response.model_dump()
+    assert serialized_response == {"message": "Resource not found"}