fix: add system v2 to generator set consumers

equinor · Sep 12, 2023 · f776b16 · f776b16
1 parent 56da4b2
commit f776b16
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/src/ecalc/libraries/libecalc/common/libecalc/input/yaml_types/components/generator_set.py b/src/ecalc/libraries/libecalc/common/libecalc/input/yaml_types/components/generator_set.py
@@ -1,11 +1,13 @@
-from typing import Any, Dict, List, Type
+from typing import Any, Dict, List, Type, Union
 
 from libecalc.dto.base import ConsumerUserDefinedCategoryType
 from libecalc.input.yaml_types import YamlBase
 from libecalc.input.yaml_types.components.category import CategoryField
+from libecalc.input.yaml_types.components.compressor_system import CompressorSystem
 from libecalc.input.yaml_types.components.legacy.yaml_electricity_consumer import (
     YamlElectricityConsumer,
 )
+from libecalc.input.yaml_types.components.pump_system import PumpSystem
 from libecalc.input.yaml_types.placeholder_type import PlaceholderType
 from libecalc.input.yaml_types.schema_helpers import (
     replace_temporal_placeholder_property_with_legacy_ref,
@@ -43,7 +45,7 @@ def schema_extra(schema: Dict[str, Any], model: Type["YamlGeneratorSet"]) -> Non
         description="Specifies the correlation between the electric power delivered and the fuel burned by a "
         "generator set.\n\n$ECALC_DOCS_KEYWORDS_URL/ELECTRICITY2FUEL",
     )
-    consumers: List[YamlElectricityConsumer] = Field(
+    consumers: List[Union[YamlElectricityConsumer, CompressorSystem, PumpSystem]] = Field(
         ...,
         title="CONSUMERS",
         description="Consumers getting electrical power from the generator set.\n\n$ECALC_DOCS_KEYWORDS_URL/CONSUMERS",

diff --git a/...t/validation/snapshots/test_validation_json_schemas/test_json_schema_changed/schemas.json b/...t/validation/snapshots/test_validation_json_schemas/test_json_schema_changed/schemas.json
@@ -408,7 +408,17 @@
                         "CONSUMERS": {
                             "description": "Consumers getting electrical power from the generator set.\n\nhttps://test.ecalc.equinor.com/docs/docs/modelling/keywords/CONSUMERS",
                             "items": {
-                                "$ref": "#/definitions/YamlElectricityConsumer"
+                                "anyOf": [
+                                    {
+                                        "$ref": "#/definitions/YamlElectricityConsumer"
+                                    },
+                                    {
+                                        "$ref": "#/definitions/CompressorSystem"
+                                    },
+                                    {
+                                        "$ref": "#/definitions/PumpSystem"
+                                    }
+                                ]
                             },
                             "title": "CONSUMERS",
                             "type": "array"