From 75a73a879ee54426e3ac8620398b3dd3769d2652 Mon Sep 17 00:00:00 2001 From: Jostein Solaas Date: Fri, 1 Nov 2024 15:08:58 +0100 Subject: [PATCH] test: introduce yaml builders Remove minimal model case --- src/libecalc/fixtures/__init__.py | 1 - .../fixtures/cases/minimal/__init__.py | 2 - .../fixtures/cases/minimal/minimal_dto.py | 66 ---- .../fixtures/cases/minimal/minimal_yaml.py | 39 --- .../yaml/yaml_models/pyyaml_yaml_model.py | 5 +- .../presentation/yaml/yaml_types/__init__.py | 4 +- src/tests/conftest.py | 108 +++++- .../integration/test_minimal_is_valid.py | 32 +- .../test_multiple_installations.py | 46 ++- src/tests/libecalc/yaml_builder.py | 327 ++++++++++++++++++ 10 files changed, 493 insertions(+), 137 deletions(-) delete mode 100644 src/libecalc/fixtures/cases/minimal/__init__.py delete mode 100644 src/libecalc/fixtures/cases/minimal/minimal_dto.py delete mode 100644 src/libecalc/fixtures/cases/minimal/minimal_yaml.py create mode 100644 src/tests/libecalc/yaml_builder.py diff --git a/src/libecalc/fixtures/__init__.py b/src/libecalc/fixtures/__init__.py index 25d157ce83..bda9c379ec 100644 --- a/src/libecalc/fixtures/__init__.py +++ b/src/libecalc/fixtures/__init__.py @@ -7,7 +7,6 @@ from .cases.consumer_with_time_slots_models import * # noqa: F403 from .cases.ltp_export import ltp_export_yaml from .cases.ltp_export.ltp_power_from_shore_yaml import ltp_pfs_yaml_factory -from .cases.minimal import * # noqa: F403 from .compressor_process_simulations.compressor_process_simulations import * # noqa: F403 from .conftest import ( fuel_gas, diff --git a/src/libecalc/fixtures/cases/minimal/__init__.py b/src/libecalc/fixtures/cases/minimal/__init__.py deleted file mode 100644 index 25c12171cb..0000000000 --- a/src/libecalc/fixtures/cases/minimal/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .minimal_dto import minimal_installation_dto_factory, minimal_model_dto_factory -from .minimal_yaml import minimal_model_yaml_factory diff --git a/src/libecalc/fixtures/cases/minimal/minimal_dto.py b/src/libecalc/fixtures/cases/minimal/minimal_dto.py deleted file mode 100644 index d01dbb63c6..0000000000 --- a/src/libecalc/fixtures/cases/minimal/minimal_dto.py +++ /dev/null @@ -1,66 +0,0 @@ -from datetime import datetime - -import pytest - -from libecalc.common.component_type import ComponentType -from libecalc.common.energy_usage_type import EnergyUsageType -from libecalc.common.time_utils import Period -from libecalc.dto import Asset, DirectConsumerFunction, Emission, FuelConsumer, FuelType, Installation -from libecalc.dto.types import ConsumerUserDefinedCategoryType, InstallationUserDefinedCategoryType -from libecalc.expression import Expression - - -def minimal_installation_dto( - installation_name: str = "minimal_installation", fuel_rate: int = 50, start: datetime = datetime(2020, 1, 1) -): - regularity = {Period(start): Expression.setup_from_expression(1)} - return Installation( - name=installation_name, - regularity=regularity, - venting_emitters=[], - hydrocarbon_export={Period(start): Expression.setup_from_expression(0)}, - user_defined_category=InstallationUserDefinedCategoryType.FIXED, - fuel_consumers=[ - FuelConsumer( - name=f"{installation_name}-direct", - component_type=ComponentType.GENERIC, - user_defined_category={Period(start): ConsumerUserDefinedCategoryType.MISCELLANEOUS}, - regularity=regularity, - fuel={ - Period(start): FuelType( - name="fuel", - emissions=[ - Emission( - name="co2", - factor=Expression.setup_from_expression(2), - ) - ], - ) - }, - energy_usage_model={ - Period(start): DirectConsumerFunction( - fuel_rate=Expression.setup_from_expression(fuel_rate), - energy_usage_type=EnergyUsageType.FUEL, - ) - }, - ) - ], - ) - - -@pytest.fixture -def minimal_installation_dto_factory(): - return minimal_installation_dto - - -@pytest.fixture -def minimal_model_dto_factory(): - def minimal_model_dto( - asset_name: str = "minimal_model", fuel_rate: int = 50, start: datetime = datetime(2020, 1, 1) - ): - return Asset( - name=asset_name, - installations=[minimal_installation_dto(fuel_rate=fuel_rate, start=start)], - ) - - return minimal_model_dto diff --git a/src/libecalc/fixtures/cases/minimal/minimal_yaml.py b/src/libecalc/fixtures/cases/minimal/minimal_yaml.py deleted file mode 100644 index db3e61049c..0000000000 --- a/src/libecalc/fixtures/cases/minimal/minimal_yaml.py +++ /dev/null @@ -1,39 +0,0 @@ -from typing import NamedTuple, Union - -import pytest - - -class YamlModel(NamedTuple): - name: str - source: str - - -@pytest.fixture -def minimal_model_yaml_factory(): - def minimal_model_yaml(fuel_rate: Union[int, str] = 50): - return YamlModel( - name="minimal_model", - source=f""" -FUEL_TYPES: - - NAME: fuel - EMISSIONS: - - NAME: co2 - FACTOR: 2 - -START: 2020-01-01 -END: 2023-01-01 - -INSTALLATIONS: - - NAME: minimal_installation - HCEXPORT: 0 - FUEL: fuel - FUELCONSUMERS: - - NAME: direct - CATEGORY: MISCELLANEOUS - ENERGY_USAGE_MODEL: - TYPE: DIRECT - FUELRATE: {fuel_rate} -""", - ) - - return minimal_model_yaml diff --git a/src/libecalc/presentation/yaml/yaml_models/pyyaml_yaml_model.py b/src/libecalc/presentation/yaml/yaml_models/pyyaml_yaml_model.py index 802f0ae437..1944c7f0ad 100644 --- a/src/libecalc/presentation/yaml/yaml_models/pyyaml_yaml_model.py +++ b/src/libecalc/presentation/yaml/yaml_models/pyyaml_yaml_model.py @@ -48,6 +48,7 @@ from libecalc.presentation.yaml.yaml_types.time_series.yaml_time_series import ( YamlTimeSeriesCollection, ) +from libecalc.presentation.yaml.yaml_types.yaml_default_datetime import YamlDefaultDatetime from libecalc.presentation.yaml.yaml_types.yaml_variable import ( YamlVariable, YamlVariableReferenceId, @@ -437,12 +438,12 @@ def installations(self) -> Iterable[YamlInstallation]: @property def start(self) -> Optional[datetime.datetime]: start_value = self._internal_datamodel.get(EcalcYamlKeywords.start) - return convert_date_to_datetime(start_value) if start_value is not None else None + return TypeAdapter(YamlDefaultDatetime).validate_python(start_value) if start_value is not None else None @property def end(self) -> Optional[datetime.datetime]: end_value = self._internal_datamodel.get(EcalcYamlKeywords.end) - return convert_date_to_datetime(end_value) if end_value is not None else None + return TypeAdapter(YamlDefaultDatetime).validate_python(end_value) if end_value is not None else None @property def dates(self): diff --git a/src/libecalc/presentation/yaml/yaml_types/__init__.py b/src/libecalc/presentation/yaml/yaml_types/__init__.py index ee3c987f74..7918fff8f4 100644 --- a/src/libecalc/presentation/yaml/yaml_types/__init__.py +++ b/src/libecalc/presentation/yaml/yaml_types/__init__.py @@ -5,5 +5,7 @@ class YamlBase(BaseModel, ABC): model_config = ConfigDict( - populate_by_name=True, alias_generator=lambda field_name: field_name.upper(), extra="forbid" + populate_by_name=True, + alias_generator=lambda field_name: field_name.upper(), + extra="forbid", ) diff --git a/src/tests/conftest.py b/src/tests/conftest.py index a1ed4b2490..f702832f73 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -1,7 +1,10 @@ import json +from io import StringIO from pathlib import Path +from typing import cast import pytest +import yaml from libecalc.common.math.numbers import Numbers from libecalc.examples import advanced, simple @@ -11,6 +14,16 @@ consumer_system_v2, ltp_export, ) +from libecalc.presentation.yaml.configuration_service import ConfigurationService +from libecalc.presentation.yaml.yaml_entities import ResourceStream +from libecalc.presentation.yaml.yaml_models.yaml_model import ReaderType, YamlConfiguration, YamlValidator +from libecalc.presentation.yaml.yaml_types.components.yaml_asset import YamlAsset +from tests.libecalc.yaml_builder import ( + YamlAssetBuilder, + YamlEnergyUsageModelDirectBuilder, + YamlFuelConsumerBuilder, + YamlInstallationBuilder, +) def _round_floats(obj): @@ -44,7 +57,6 @@ def rounded_snapshot(data: dict, snapshot_name: str): "consumer_system_v2": (Path(consumer_system_v2.__file__).parent / "data" / "consumer_system_v2.yaml").absolute(), } - # The value should be the name of a fixture returning the YamlCase for the example valid_example_yaml_case_fixture_names = { "simple": "simple_yaml", @@ -110,3 +122,97 @@ def valid_example_case_yaml_case(request) -> YamlCase: """ yaml_case = request.getfixturevalue(request.param[1]) return yaml_case + + +class YamlAssetConfigurationService(ConfigurationService): + def __init__(self, model: YamlAsset, name: str): + self._name = name + self._model = model + self._source = None + + @property + def name(self): + return self._name + + @property + def source(self): + return self.get_yaml() + + def get_yaml(self) -> str: + if self._source is not None: + return self._source + data = self._model.model_dump(by_alias=True, exclude_unset=True, mode="json") + return yaml.dump(data) + + def get_configuration(self) -> YamlValidator: + stream = ResourceStream(name=self._name, stream=StringIO(self.get_yaml())) + main_yaml_model = YamlConfiguration.Builder.get_yaml_reader(ReaderType.PYYAML).read( + main_yaml=stream, + enable_include=True, + ) + return cast(YamlValidator, main_yaml_model) + + +@pytest.fixture +def yaml_asset_configuration_service_factory(): + def yaml_asset_configuration_service(model: YamlAsset, name: str): + return YamlAssetConfigurationService(model=model, name=name) + + return yaml_asset_configuration_service + + +@pytest.fixture +def yaml_asset_builder_factory(): + return lambda: YamlAssetBuilder() + + +@pytest.fixture +def yaml_installation_builder_factory(): + return lambda: YamlInstallationBuilder() + + +@pytest.fixture +def minimal_installation_yaml_factory(yaml_installation_builder_factory): + def minimal_installation_yaml( + name: str = "DefaultInstallation", + consumer_name: str = "flare", + fuel_name: str = "fuel", + fuel_rate: int | str = 50, + ): + return ( + yaml_installation_builder_factory() + .with_test_data() + .with_name(name) + .with_fuel_consumers( + [ + YamlFuelConsumerBuilder() + .with_test_data() + .with_name(consumer_name) + .with_fuel(fuel_name) + .with_energy_usage_model( + YamlEnergyUsageModelDirectBuilder().with_test_data().with_fuel_rate(fuel_rate).build() + ) + .build() + ] + ) + .build() + ) + + return minimal_installation_yaml + + +@pytest.fixture +def minimal_model_yaml_factory(minimal_installation_yaml_factory): + def minimal_model_yaml(fuel_rate: int | str = 50): + fuel_name = "fuel" + installation = minimal_installation_yaml_factory(fuel_name="fuel", fuel_rate=fuel_rate) + model = ( + YamlAssetBuilder() + .with_test_data(fuel_name=fuel_name) + .with_installations([installation]) + .with_start("2020-01-01") + .with_end("2023-01-01") + ) + return YamlAssetConfigurationService(model.build(), name="minimal_model") + + return minimal_model_yaml diff --git a/src/tests/libecalc/integration/test_minimal_is_valid.py b/src/tests/libecalc/integration/test_minimal_is_valid.py index 4c173fe76d..435a38e4bd 100644 --- a/src/tests/libecalc/integration/test_minimal_is_valid.py +++ b/src/tests/libecalc/integration/test_minimal_is_valid.py @@ -1,33 +1,41 @@ -from datetime import datetime - import pytest from libecalc.application.energy_calculator import EnergyCalculator from libecalc.application.graph_result import GraphResult -from libecalc.common.variables import VariablesMap +from libecalc.common.time_utils import Frequency from libecalc.presentation.json_result.mapper import get_asset_result from libecalc.presentation.json_result.result import EcalcModelResult +from libecalc.presentation.yaml.model import YamlModel +from libecalc.presentation.yaml.resource import Resource +from libecalc.presentation.yaml.resource_service import ResourceService +from libecalc.presentation.yaml.yaml_models.yaml_model import YamlValidator + + +class EmptyResourceService(ResourceService): + def get_resources(self, configuration: YamlValidator) -> dict[str, Resource]: + return {} @pytest.fixture -def minimal_asset_result(minimal_model_dto_factory): - minimal_dto = minimal_model_dto_factory() - graph = minimal_dto.get_graph() - variables = VariablesMap( - time_vector=[datetime(2020, 1, 1), datetime(2022, 1, 1)], - variables={}, +def minimal_asset_result(minimal_model_yaml_factory): + minimal_configuration_service = minimal_model_yaml_factory() + model = YamlModel( + configuration_service=minimal_configuration_service, + resource_service=EmptyResourceService(), + output_frequency=Frequency.NONE, ) + graph = model.get_graph() energy_calculator = EnergyCalculator(graph=graph) - consumer_results = energy_calculator.evaluate_energy_usage(variables) + consumer_results = energy_calculator.evaluate_energy_usage(model.variables) emission_results = energy_calculator.evaluate_emissions( - variables_map=variables, + variables_map=model.variables, consumer_results=consumer_results, ) return get_asset_result( GraphResult( graph=graph, consumer_results=consumer_results, - variables_map=variables, + variables_map=model.variables, emission_results=emission_results, ) ) diff --git a/src/tests/libecalc/integration/test_multiple_installations.py b/src/tests/libecalc/integration/test_multiple_installations.py index 588bff4bc5..9d568a55da 100644 --- a/src/tests/libecalc/integration/test_multiple_installations.py +++ b/src/tests/libecalc/integration/test_multiple_installations.py @@ -2,31 +2,51 @@ import pytest -from libecalc import dto from libecalc.application.energy_calculator import EnergyCalculator from libecalc.application.graph_result import GraphResult +from libecalc.common.time_utils import Frequency from libecalc.common.units import Unit from libecalc.common.utils.rates import RateType, TimeSeriesRate from libecalc.common.variables import VariablesMap from libecalc.presentation.json_result.mapper import get_asset_result +from libecalc.presentation.yaml.model import YamlModel +from libecalc.presentation.yaml.resource import Resource +from libecalc.presentation.yaml.resource_service import ResourceService +from libecalc.presentation.yaml.yaml_models.yaml_model import YamlValidator + + +class EmptyResourceService(ResourceService): + def get_resources(self, configuration: YamlValidator) -> dict[str, Resource]: + return {} @pytest.fixture -def asset_with_two_installations(minimal_installation_dto_factory) -> dto.Asset: - installation_1 = minimal_installation_dto_factory(installation_name="installaion1", fuel_rate=50) - installation_2 = minimal_installation_dto_factory(installation_name="installaion2", fuel_rate=100) - asset = dto.Asset( - name="multiple_installations_asset", - installations=[ - installation_1, - installation_2, - ], +def model_with_two_installations( + minimal_installation_yaml_factory, yaml_asset_configuration_service_factory, yaml_asset_builder_factory +) -> YamlModel: + installation_1 = minimal_installation_yaml_factory( + name="installation1", fuel_rate=50, fuel_name="fuel", consumer_name="flare1" + ) + installation_2 = minimal_installation_yaml_factory( + name="installation2", fuel_rate=100, fuel_name="fuel", consumer_name="flare2" + ) + + asset = ( + yaml_asset_builder_factory() + .with_test_data(fuel_name="fuel") + .with_installations([installation_1, installation_2]) + .build() + ) + + return YamlModel( + configuration_service=yaml_asset_configuration_service_factory(asset, "multiple_installations_asset"), + resource_service=EmptyResourceService(), + output_frequency=Frequency.YEAR, ) - return asset -def test_asset_with_multiple_installations(asset_with_two_installations): - graph = asset_with_two_installations.get_graph() +def test_asset_with_multiple_installations(model_with_two_installations): + graph = model_with_two_installations.get_graph() energy_calculator = EnergyCalculator(graph) timesteps = [ datetime(2020, 1, 1), diff --git a/src/tests/libecalc/yaml_builder.py b/src/tests/libecalc/yaml_builder.py new file mode 100644 index 0000000000..948c9d33a7 --- /dev/null +++ b/src/tests/libecalc/yaml_builder.py @@ -0,0 +1,327 @@ +import abc +from enum import Enum +from typing import List, Self, TypeVar, Generic, Any + +from pydantic import TypeAdapter + +from libecalc.dto.types import ConsumerUserDefinedCategoryType +from libecalc.presentation.yaml.yaml_types import YamlBase +from libecalc.presentation.yaml.yaml_types.components.legacy.energy_usage_model import YamlFuelEnergyUsageModel +from libecalc.presentation.yaml.yaml_types.components.legacy.energy_usage_model.yaml_energy_usage_model_direct import ( + ConsumptionRateType, + YamlEnergyUsageModelDirect, +) +from libecalc.presentation.yaml.yaml_types.components.legacy.yaml_fuel_consumer import YamlFuelConsumer +from libecalc.presentation.yaml.yaml_types.components.yaml_asset import YamlAsset +from libecalc.presentation.yaml.yaml_types.components.yaml_expression_type import YamlExpressionType +from libecalc.presentation.yaml.yaml_types.components.yaml_installation import YamlInstallation +from libecalc.presentation.yaml.yaml_types.facility_model.yaml_facility_model import ( + YamlFacilityModel, +) +from libecalc.presentation.yaml.yaml_types.fuel_type.yaml_emission import YamlEmission +from libecalc.presentation.yaml.yaml_types.fuel_type.yaml_fuel_type import YamlFuelType +from libecalc.presentation.yaml.yaml_types.models import YamlConsumerModel +from libecalc.presentation.yaml.yaml_types.time_series.yaml_time_series import ( + YamlTimeSeriesCollection, +) +from libecalc.presentation.yaml.yaml_types.yaml_temporal_model import YamlTemporalModel +from libecalc.presentation.yaml.yaml_types.yaml_variable import YamlVariables + +T = TypeVar("T") + + +class Builder(abc.ABC, Generic[T]): + @abc.abstractmethod + def with_test_data(self) -> Self: ... + + @abc.abstractmethod + def build(self) -> T: ... + + +TClass = TypeVar("TClass", bound=YamlBase) + + +def _build(obj: Any, attrs: list[str], c: type[TClass], exclude_none=True): + data = {} + for attr in attrs: + value = getattr(obj, attr) + if isinstance(value, Enum): + value = value.value + + if value is None: + continue + data[attr] = value + + return TypeAdapter(c).validate_python(data) + + +class YamlEnergyUsageModelDirectBuilder: + def __init__(self): + self.type = "DIRECT" + self.load = None # To be set with test data or custom value + self.fuel_rate = None # To be set with test data or custom value + self.consumption_rate_type = None # To be set with test data or custom value + + def with_test_data(self): + self.load = None + self.fuel_rate = "0.5" + self.consumption_rate_type = ConsumptionRateType.STREAM_DAY.value + return self + + def with_load(self, load: YamlExpressionType): + self.load = load + return self + + def with_fuel_rate(self, fuel_rate: YamlExpressionType): + self.fuel_rate = fuel_rate + return self + + def with_consumption_rate_type(self, consumption_rate_type: ConsumptionRateType): + self.consumption_rate_type = consumption_rate_type.value + return self + + def build(self) -> YamlEnergyUsageModelDirect: + return _build(self, ["type", "load", "fuel_rate", "consumption_rate_type"], YamlEnergyUsageModelDirect) + + +TYamlClass = TypeVar("TYamlClass", bound=YamlBase) + + +class YamlModelContainer(Generic[TYamlClass]): + """ + Consumer, Installation wrapper that can be used to bundle the component with its dependencies. + + A fuel consumer that uses a model reference can then use this class to keep the referenced model in a single class. + """ + + def __init__(self, component: TYamlClass): + self.component = component + self.models = None + self.time_series = None + self.facility_inputs = None + + def with_models(self, models) -> Self: + self.models = models + return self + + def with_time_series(self, time_series): + self.time_series = time_series + return self + + def with_facility_inputs(self, facility_inputs) -> Self: + self.facility_inputs = facility_inputs + return self + + +class YamlFuelConsumerBuilder: + def __init__(self): + self.name = None + self.fuel = None + self.energy_usage_model = None + self.category = None + + def with_test_data(self) -> Self: + self.name = "flare" + self.category = ConsumerUserDefinedCategoryType.FLARE.value + self.energy_usage_model = YamlEnergyUsageModelDirectBuilder().with_test_data().build() + self.fuel = None + return self + + def with_name(self, name: str) -> Self: + self.name = name + return self + + def with_fuel(self, fuel: str) -> Self: + self.fuel = fuel + return self + + def with_energy_usage_model(self, energy_usage_model: YamlTemporalModel[YamlFuelEnergyUsageModel]) -> Self: + self.energy_usage_model = energy_usage_model + return self + + def build(self) -> YamlFuelConsumer: + return _build(self, ["name", "category", "fuel", "energy_usage_model"], YamlFuelConsumer) + + +class YamlInstallationBuilder: + def __init__(self): + self.name = None + self.category = None # Placeholder for InstallationUserDefinedCategoryType + self.hydrocarbon_export = None # Placeholder for YamlTemporalModel[YamlExpressionType] + self.fuel = None # Placeholder for YamlTemporalModel[str] + self.regularity = None # Placeholder for YamlTemporalModel[YamlExpressionType] + self.generator_sets = [] # Placeholder for List[YamlGeneratorSet] + self.fuel_consumers = [] # Placeholder for List[Union[YamlFuelConsumer, YamlConsumerSystem]] + self.venting_emitters = [] # Placeholder for List[YamlVentingEmitter] + + self._yaml_model_containers = [] + + def with_test_data(self): + # Populate with test data if needed + # For example, create a list of YamlGeneratorSet, YamlFuelConsumer, YamlConsumerSystem, and YamlVentingEmitter instances + # self.generator_sets.append(YamlGeneratorSetBuilder().with_test_data().build()) + # self.fuel_consumers.append(YamlFuelConsumerBuilder().with_test_data().build()) + # self.venting_emitters.append(YamlVentingEmitterBuilder().with_test_data().build()) + self.name = "DefaultInstallation" + self.hydrocarbon_export = 0 + self.regularity = 1 + self.fuel_consumers.append(YamlFuelConsumerBuilder().with_test_data().build()) + return self + + def with_name(self, name: str) -> Self: + self.name = name + return self + + def with_fuel_consumers( + self, fuel_consumers: list[YamlFuelConsumer | YamlModelContainer[YamlFuelConsumer]] + ) -> Self: + new_fuel_consumers = [] + for fuel_consumer in fuel_consumers: + if isinstance(fuel_consumer, YamlModelContainer): + self._yaml_model_containers.append(fuel_consumer) + new_fuel_consumers.append(fuel_consumer.component) + else: + new_fuel_consumers.append(fuel_consumer) + self.fuel_consumers = new_fuel_consumers + return self + + def build(self) -> YamlInstallation: + return _build( + self, + [ + "name", + "category", + "fuel", + "hydrocarbon_export", + "regularity", + "generator_sets", + "fuel_consumers", + "venting_emitters", + ], + YamlInstallation, + ) + + +class YamlEmissionBuilder: + def __init__(self): + self.name = None + self.factor = None + + def with_test_data(self) -> Self: + # Populate with test data if needed + self.name = "CO2" + self.factor = 2 + return self + + def build(self) -> YamlEmission: + return _build(self, ["name", "factor"], YamlEmission) + + +# Builder for YamlFuelType +class YamlFuelTypeBuilder: + def __init__(self): + self.name = None + self.category = None # Placeholder for FuelTypeUserDefinedCategoryType + self.emissions = [] # Placeholder for a list of YamlEmission + + def with_test_data(self) -> Self: + # Populate with test data if needed + # For example, create a list of YamlEmission instances using YamlEmissionBuilder + self.name = "fuel" + emission_builder = YamlEmissionBuilder().with_test_data() + self.emissions.append(emission_builder.build()) + return self + + def with_name(self, name: str) -> Self: + self.name = name + return self + + def build(self) -> YamlFuelType: + return _build(self, ["name", "category", "emissions"], YamlFuelType) + + +class YamlAssetBuilder: + def __init__(self): + self.time_series = [] + self.facility_inputs = [] + self.models = [] + self.fuel_types = [] + self.variables = None + self.installations = [] + self.start = None + self.end = None + + def with_test_data(self, fuel_name: str = "fuel"): + self.time_series = [] + + self.facility_inputs = [] + + self.models = [] + + # Create and populate a list of YamlFuelType test instances + fuel_types_builder = YamlFuelTypeBuilder().with_test_data().with_name(fuel_name) + self.fuel_types = [fuel_types_builder.build()] + + self.variables = None + + # Create and populate a list of YamlInstallation test instances + installations_builder = YamlInstallationBuilder().with_test_data() + self.installations = [installations_builder.build()] + + self.start = "2019-01-01" + self.end = "2024-01-01" + + return self + + def with_time_series(self, time_series: List[YamlTimeSeriesCollection]): + self.time_series = time_series + return self + + def with_facility_inputs(self, facility_inputs: List[YamlFacilityModel]): + self.facility_inputs = facility_inputs + return self + + def with_models(self, models: List[YamlConsumerModel]): + self.models = models + return self + + def with_fuel_types(self, fuel_types: List[YamlFuelType]): + self.fuel_types = fuel_types + return self + + def with_variables(self, variables: YamlVariables): + self.variables = variables + return self + + def with_installations(self, installations: List[YamlInstallation | YamlModelContainer[YamlInstallation]]): + if len(installations) == 0: + self.installations = [] + return self + + new_installations = [] + for installation in installations: + if isinstance(installation, YamlModelContainer): + self.models.extend(installation.models) + self.facility_inputs.extend(installation.facility_inputs) + self.time_series.extend(installation.time_series) + new_installations.append(installation.component) + else: + new_installations.append(installation) + + self.installations = new_installations + return self + + def with_start(self, start: str): + self.start = start + return self + + def with_end(self, end: str): + self.end = end + return self + + def build(self) -> YamlAsset: + return _build( + self, + ["time_series", "facility_inputs", "models", "fuel_types", "variables", "installations", "start", "end"], + YamlAsset, + )