From f776b16df627ea08d802356ea63a4ce077f45094 Mon Sep 17 00:00:00 2001
From: Jostein Solaas <josteinsolaas@gmail.com>
Date: Tue, 12 Sep 2023 11:34:45 +0200
Subject: [PATCH] fix: add system v2 to generator set consumers

---
 .../input/yaml_types/components/generator_set.py     |  6 ++++--
 .../test_json_schema_changed/schemas.json            | 12 +++++++++++-
 2 files changed, 15 insertions(+), 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
index b13f9c2d3e..775cb91a00 100644
--- 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/src/ecalc/libraries/libecalc/common/tests/input/validation/snapshots/test_validation_json_schemas/test_json_schema_changed/schemas.json b/src/ecalc/libraries/libecalc/common/tests/input/validation/snapshots/test_validation_json_schemas/test_json_schema_changed/schemas.json
index 9ab0e431e0..27677cb65c 100644
--- a/src/ecalc/libraries/libecalc/common/tests/input/validation/snapshots/test_validation_json_schemas/test_json_schema_changed/schemas.json
+++ b/src/ecalc/libraries/libecalc/common/tests/input/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"