Skip to content

Commit

Permalink
feat(load): add area integration test
Browse files Browse the repository at this point in the history
  • Loading branch information
@TheoPascoli
TheoPascoli committed Jan 30, 2025
1 parent c9fcd01 commit f8f1570
Show file tree
Hide file tree
Showing 3 changed files with 128 additions and 84 deletions.
52 changes: 48 additions & 4 deletions antarest/study/business/area_management.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
# This file is part of the Antares project.

import logging
import re
import typing as t

from antarest.core.exceptions import ConfigFileNotFound, DuplicateAreaName, LayerNotAllowedToBeDeleted, LayerNotFound
Expand All @@ -23,14 +24,16 @@
ClusterInfoDTO,
LayerInfoDTO,
UpdateAreaUi,
_get_area_layers,
_get_ui_info_map,
)
from antarest.study.business.utils import execute_or_add_commands
from antarest.study.model import Patch, PatchArea, PatchCluster, RawStudy, Study
from antarest.study.repository import StudyMetadataRepository
from antarest.study.storage.patch_service import PatchService
from antarest.study.storage.rawstudy.model.filesystem.config.area import AreaFolder, ThermalAreasProperties
from antarest.study.storage.rawstudy.model.filesystem.config.area import (
AreaFolder,
ThermalAreasProperties,
UIProperties,
)
from antarest.study.storage.rawstudy.model.filesystem.config.model import Area, DistrictSet, transform_name_to_id
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
from antarest.study.storage.storage_service import StudyStorageService
Expand All @@ -47,6 +50,47 @@
_THERMAL_AREAS_PATH = "input/thermal/areas"


def _get_ui_info_map(file_study: FileStudy, area_ids: t.Sequence[str]) -> t.Dict[str, t.Any]:
"""
Get the UI information (a JSON object) for each selected Area.
Args:
file_study: A file study from which the configuration can be read.
area_ids: List of selected area IDs.
Returns:
Dictionary where keys are IDs, and values are UI objects.
Raises:
ChildNotFoundError: if one of the Area IDs is not found in the configuration.
"""
# If there is no ID, it is better to return an empty dictionary
# instead of raising an obscure exception.
if not area_ids:
return {}

ui_info_map = file_study.tree.get(["input", "areas", ",".join(area_ids), "ui"])

# If there is only one ID in the `area_ids`, the result returned from
# the `file_study.tree.get` call will be a single UI object.
# On the other hand, if there are multiple values in `area_ids`,
# the result will be a dictionary where the keys are the IDs,
# and the values are the corresponding UI objects.
if len(area_ids) == 1:
ui_info_map = {area_ids[0]: ui_info_map}

# Convert to UIProperties to ensure that the UI object is valid.
ui_info_map = {area_id: UIProperties(**ui_info).to_config() for area_id, ui_info in ui_info_map.items()}

return ui_info_map


def _get_area_layers(area_uis: t.Dict[str, t.Any], area: str) -> t.List[str]:
if area in area_uis and "ui" in area_uis[area] and "layers" in area_uis[area]["ui"]:
return re.split(r"\s+", (str(area_uis[area]["ui"]["layers"]) or ""))
return []


class AreaManager:
"""
Manages operations related to areas in a study, including retrieval, creation, and updates.
Expand Down Expand Up @@ -195,7 +239,7 @@ def update_areas_props(

@staticmethod
def get_table_schema() -> JSON:
return AreaOutput.schema()
return AreaOutput.model_json_schema()

def get_all_areas(self, study: RawStudy, area_type: t.Optional[AreaType] = None) -> t.List[AreaInfoDTO]:
"""
Expand Down
81 changes: 1 addition & 80 deletions antarest/study/business/model/area_model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
# SPDX-License-Identifier: MPL-2.0
#
# This file is part of the Antares project.

True
import enum
import re
import typing as t
Expand Down Expand Up @@ -69,44 +69,6 @@ class LayerInfoDTO(AntaresBaseModel):


class UpdateAreaUi(AntaresBaseModel, extra="forbid", populate_by_name=True):
"""
DTO for updating area UI
Usage:
>>> from antarest.study.business.model.area_model import UpdateAreaUi
>>> from pprint import pprint
>>> obj = {
... "x": -673.75,
... "y": 301.5,
... "color_rgb": [230, 108, 44],
... "layerX": {"0": -230, "4": -230, "6": -95, "7": -230, "8": -230},
... "layerY": {"0": 136, "4": 136, "6": 39, "7": 136, "8": 136},
... "layerColor": {
... "0": "230, 108, 44",
... "4": "230, 108, 44",
... "6": "230, 108, 44",
... "7": "230, 108, 44",
... "8": "230, 108, 44",
... },
... }
>>> model = UpdateAreaUi(**obj)
>>> pprint(model.model_dump(by_alias=True), width=80)
{'colorRgb': [230, 108, 44],
'layerColor': {0: '230, 108, 44',
4: '230, 108, 44',
6: '230, 108, 44',
7: '230, 108, 44',
8: '230, 108, 44'},
'layerX': {0: -230, 4: -230, 6: -95, 7: -230, 8: -230},
'layerY': {0: 136, 4: 136, 6: 39, 7: 136, 8: 136},
'x': -673,
'y': 301}
"""

x: int = Field(title="X position")
y: int = Field(title="Y position")
color_rgb: t.Sequence[int] = Field(title="RGB color", alias="colorRgb")
Expand All @@ -115,47 +77,6 @@ class UpdateAreaUi(AntaresBaseModel, extra="forbid", populate_by_name=True):
layer_color: t.Mapping[int, str] = Field(default_factory=dict, title="Color of each layer", alias="layerColor")


def _get_ui_info_map(file_study: FileStudy, area_ids: t.Sequence[str]) -> t.Dict[str, t.Any]:
"""
Get the UI information (a JSON object) for each selected Area.
Args:
file_study: A file study from which the configuration can be read.
area_ids: List of selected area IDs.
Returns:
Dictionary where keys are IDs, and values are UI objects.
Raises:
ChildNotFoundError: if one of the Area IDs is not found in the configuration.
"""
# If there is no ID, it is better to return an empty dictionary
# instead of raising an obscure exception.
if not area_ids:
return {}

ui_info_map = file_study.tree.get(["input", "areas", ",".join(area_ids), "ui"])

# If there is only one ID in the `area_ids`, the result returned from
# the `file_study.tree.get` call will be a single UI object.
# On the other hand, if there are multiple values in `area_ids`,
# the result will be a dictionary where the keys are the IDs,
# and the values are the corresponding UI objects.
if len(area_ids) == 1:
ui_info_map = {area_ids[0]: ui_info_map}

# Convert to UIProperties to ensure that the UI object is valid.
ui_info_map = {area_id: UIProperties(**ui_info).to_config() for area_id, ui_info in ui_info_map.items()}

return ui_info_map


def _get_area_layers(area_uis: t.Dict[str, t.Any], area: str) -> t.List[str]:
if area in area_uis and "ui" in area_uis[area] and "layers" in area_uis[area]["ui"]:
return re.split(r"\s+", (str(area_uis[area]["ui"]["layers"]) or ""))
return []


# noinspection SpellCheckingInspection
class _BaseAreaDTO(
OptimizationProperties.FilteringSection,
Expand Down
79 changes: 79 additions & 0 deletions tests/integration/study_data_blueprint/test_area.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
# Copyright (c) 2025, RTE (https://www.rte-france.com)
#
# See AUTHORS.txt
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
#
# SPDX-License-Identifier: MPL-2.0
#
# This file is part of the Antares project.

import pytest
from starlette.testclient import TestClient

from tests.integration.prepare_proxy import PreparerProxy


@pytest.mark.unit_test
class TestArea:
@pytest.mark.parametrize("study_type", ["raw", "variant"])
def test_area(self, client: TestClient, user_access_token: str, study_type: str) -> None:
client.headers = {"Authorization": f"Bearer {user_access_token}"} # type: ignore

preparer = PreparerProxy(client, user_access_token)
study_id = preparer.create_study("foo", version=820)
if study_type == "variant":
study_id = preparer.create_variant(study_id, name="Variant 1")

client.post(f"/v1/studies/{study_id}/areas", json={"name": "area1", "type": "AREA"})
client.post(f"/v1/studies/{study_id}/areas", json={"name": "area2", "type": "AREA"})

res = client.get(f"/v1/studies/{study_id}/areas?ui=true")
assert res.status_code == 200
expected = {
"area1": {
"layerColor": {"0": "230, 108, 44"},
"layerX": {"0": 0},
"layerY": {"0": 0},
"ui": {"color_b": 44, "color_g": 108, "color_r": 230, "layers": "0", "x": 0, "y": 0},
},
"area2": {
"layerColor": {"0": "230, 108, 44"},
"layerX": {"0": 0},
"layerY": {"0": 0},
"ui": {"color_b": 44, "color_g": 108, "color_r": 230, "layers": "0", "x": 0, "y": 0},
},
}
assert res.json() == expected

client.put(
f"/v1/studies/{study_id}/areas/area1/ui",
json={
"x": 10,
"y": 10,
"layerColor": {"0": "100, 100, 100"},
"layerX": {"0": 10},
"layerY": {"0": 10},
"color_rgb": (100, 100, 100),
},
)

res = client.get(f"/v1/studies/{study_id}/areas?ui=true")
assert res.status_code == 200
expected = {
"area1": {
"layerColor": {"0": "100, 100, 100"},
"layerX": {"0": 10},
"layerY": {"0": 10},
"ui": {"color_b": 100, "color_g": 100, "color_r": 100, "layers": "0", "x": 10, "y": 10},
},
"area2": {
"layerColor": {"0": "230, 108, 44"},
"layerX": {"0": 0},
"layerY": {"0": 0},
"ui": {"color_b": 44, "color_g": 108, "color_r": 230, "layers": "0", "x": 0, "y": 0},
},
}
assert res.json() == expected

0 comments on commit f8f1570

Please sign in to comment.