From ad99ebfbe9b83917ad079e7ee9e28bd761c79576 Mon Sep 17 00:00:00 2001
From: Stefan Verhoeven <s.verhoeven@esciencecenter.nl>
Date: Thu, 5 Oct 2023 10:01:48 +0200
Subject: [PATCH] Tests for eWaterCycleModel (#378)

* Replace test_abstract with base/test_model

* Add more tests for base.model module

* Dont depend on fake models from grpc4bmi

* Make type of eWaterCycleModel.parameters property an ItemsView

To be more inline with pymt

See https://github.com/eWaterCycle/ewatercycle/pull/365#discussion_r1328434215

* Use local var instead of prop

* Fix more tests

* More and better tests

* Use public interface where possible

* Update model.py

* Sort imports
---
 CHANGELOG.md                               |   1 +
 src/ewatercycle/base/model.py              |  32 +-
 src/ewatercycle/plugins/hype/model.py      |   6 +-
 src/ewatercycle/plugins/lisflood/model.py  |   6 +-
 src/ewatercycle/plugins/marrmot/model.py   |  10 +-
 src/ewatercycle/plugins/pcrglobwb/model.py |   6 +-
 src/ewatercycle/plugins/wflow/model.py     |   6 +-
 src/ewatercycle/testing/fake_models.py     | 134 +++++++
 tests/plugins/hype/test_model.py           |  24 +-
 tests/plugins/lisflood/test_model.py       |  27 +-
 tests/plugins/marrmot/test_model_m01.py    |   4 +-
 tests/plugins/marrmot/test_model_m14.py    |   4 +-
 tests/plugins/wflow/test_model.py          |  17 +-
 tests/src/base/__init__.py                 |   1 +
 tests/src/base/test_model.py               | 397 +++++++++++++++++++
 tests/src/fake_models.py                   | 432 ---------------------
 tests/src/models/test_abstract.py          | 339 ----------------
 tests/src/test_container.py                |  38 +-
 18 files changed, 633 insertions(+), 851 deletions(-)
 create mode 100644 tests/src/base/__init__.py
 create mode 100644 tests/src/base/test_model.py
 delete mode 100644 tests/src/fake_models.py
 delete mode 100644 tests/src/models/test_abstract.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3db9779e..c48b5e43 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,6 +26,7 @@ Formatted as described on [https://keepachangelog.com](https://keepachangelog.co
 - Forcing ((#365)[https://github.com/eWaterCycle/ewatercycle/pull/365]):
   - Instead of modifying an existing recipe now builds a ESMValTool recipe from scratch using a fluent interface
   - DefaultForcing has overridable class methods for each step of the forcing generation process (build_recipe, run_recipe, recipe_output_to_forcing_arguments).
+- eWaterCycleModel.parameters property type is ItemsView instead of dict.
 
 ### Deprecated
 
diff --git a/src/ewatercycle/base/model.py b/src/ewatercycle/base/model.py
index a2c62bb4..35a38609 100644
--- a/src/ewatercycle/base/model.py
+++ b/src/ewatercycle/base/model.py
@@ -4,6 +4,7 @@
 import datetime
 import inspect
 import logging
+from collections.abc import ItemsView
 from datetime import timezone
 from pathlib import Path
 from typing import Annotated, Any, Iterable, Optional, Type, cast
@@ -14,7 +15,6 @@
 import yaml
 from cftime import num2pydate
 from grpc4bmi.bmi_optionaldest import OptionalDestBmi
-from grpc4bmi.reserve import reserve_values, reserve_values_at_indices
 from pydantic import (
     BaseModel,
     BeforeValidator,
@@ -88,9 +88,9 @@ def _make_bmi_instance(self) -> OptionalDestBmi:
     # where it is {}.items()
     # TODO is this OK?
     @property
-    def parameters(self) -> dict[str, Any]:
+    def parameters(self) -> ItemsView[str, Any]:
         """Display the model's parameters and their values."""
-        return {}
+        return {}.items()
 
     def setup(self, *, cfg_dir: str | None = None, **kwargs) -> tuple[str, str]:
         """Perform model setup.
@@ -136,9 +136,10 @@ def _make_cfg_dir(self, cfg_dir: Optional[str] = None, **kwargs) -> Path:
     def _make_cfg_file(self, **kwargs):
         """Create new config file and return its path."""
         cfg_file = self._cfg_dir / "config.yaml"
-        self.parameters.update(**kwargs)
+        myparameters = dict(list(self.parameters))
+        myparameters.update(**kwargs)
         with cfg_file.open(mode="w") as file:
-            yaml.dump({k: v for k, v in self.parameters}, file)
+            yaml.dump({k: v for k, v in myparameters}, file)
 
         return cfg_file
 
@@ -167,7 +168,10 @@ def initialize(self, config_file: str) -> None:
         self._bmi.initialize(config_file)
 
     def finalize(self) -> None:
-        """Perform tear-down tasks for the model."""
+        """Perform tear-down tasks for the model.
+
+        After finalization, the model should not be used anymore.
+        """
         self._bmi.finalize()
         del self._bmi
 
@@ -181,10 +185,7 @@ def get_value(self, name: str) -> np.ndarray:
         Args:
             name: Name of variable
         """
-        if isinstance(self._bmi, OptionalDestBmi):
-            return self._bmi.get_value(name)
-        dest = reserve_values(self._bmi, name)
-        return self._bmi.get_value(name, dest)
+        return self._bmi.get_value(name)
 
     def get_value_at_coords(
         self, name, lat: Iterable[float], lon: Iterable[float]
@@ -198,10 +199,7 @@ def get_value_at_coords(
         """
         indices = self._coords_to_indices(name, lat, lon)
         indices = np.array(indices)
-        if isinstance(self._bmi, OptionalDestBmi):
-            return self._bmi.get_value_at_indices(name, indices)
-        dest = reserve_values_at_indices(self._bmi, name, indices)
-        return self._bmi.get_value_at_indices(name, dest, indices)
+        return self._bmi.get_value_at_indices(name, indices)
 
     def set_value(self, name: str, value: np.ndarray) -> None:
         """Specify a new value for a model variable.
@@ -273,9 +271,9 @@ def get_value_as_xarray(self, name: str) -> xr.DataArray:
             data=np.reshape(
                 self.get_value(name),
                 (
+                    1,
                     shape[0],
                     shape[1],
-                    1,
                 ),
             ),
             coords={
@@ -371,6 +369,10 @@ def end_time_as_datetime(self) -> datetime.datetime:
     @property
     def time_as_datetime(self) -> datetime.datetime:
         """Current time of the model as a datetime object'."""
+        # TODO some bmi implementations like Wflow.jl returns 'd'
+        # which can not be converted to a datetime object
+        # as nupmy2date expects a
+        # `<time units> since <reference time>` formatted string
         return num2pydate(
             self._bmi.get_current_time(),
             self._bmi.get_time_units(),
diff --git a/src/ewatercycle/plugins/hype/model.py b/src/ewatercycle/plugins/hype/model.py
index 43ee7535..380715aa 100644
--- a/src/ewatercycle/plugins/hype/model.py
+++ b/src/ewatercycle/plugins/hype/model.py
@@ -3,7 +3,7 @@
 import logging
 import types
 from pathlib import Path
-from typing import Any, Iterable, Optional
+from typing import Any, ItemsView, Iterable, Optional
 
 import bmipy
 import xarray as xr
@@ -133,7 +133,7 @@ def get_latlon_grid(self, name: str) -> tuple[Any, Any, Any]:
         raise NotImplementedError("Hype coordinates cannot be mapped to grid")
 
     @property
-    def parameters(self) -> dict[str, Any]:
+    def parameters(self) -> ItemsView[str, Any]:
         """List the parameters for this model.
 
         Exposed Lisflood parameters:
@@ -157,7 +157,7 @@ def parameters(self) -> dict[str, Any]:
             "crit_time": _get_hype_time(
                 _get_code_in_cfg(self._config, "cdate")
             ).strftime(ISO_TIMEFMT),
-        }
+        }.items()
 
     def _coords_to_indices(
         self, name: str, lat: Iterable[float], lon: Iterable[float]
diff --git a/src/ewatercycle/plugins/lisflood/model.py b/src/ewatercycle/plugins/lisflood/model.py
index 975ab5f8..aa039677 100644
--- a/src/ewatercycle/plugins/lisflood/model.py
+++ b/src/ewatercycle/plugins/lisflood/model.py
@@ -2,7 +2,7 @@
 
 import logging
 from pathlib import Path
-from typing import Any
+from typing import Any, ItemsView
 
 from pydantic import PrivateAttr, model_validator
 
@@ -137,7 +137,7 @@ def _make_cfg_file(self, **kwargs) -> Path:
         return lisflood_file
 
     @property
-    def parameters(self) -> dict[str, Any]:
+    def parameters(self) -> ItemsView[str, Any]:
         """List the parameters for this model.
 
         Exposed Lisflood parameters:
@@ -158,7 +158,7 @@ def parameters(self) -> dict[str, Any]:
             "MaskMap": self._get_textvar_value("MaskMap"),
             "start_time": get_time(self.forcing.start_time).strftime(ISO_TIMEFMT),
             "end_time": get_time(self.forcing.end_time).strftime(ISO_TIMEFMT),
-        }
+        }.items()
 
     def finalize(self) -> None:
         """Perform tear-down tasks for the model."""
diff --git a/src/ewatercycle/plugins/marrmot/model.py b/src/ewatercycle/plugins/marrmot/model.py
index a6252a27..fa42eb62 100644
--- a/src/ewatercycle/plugins/marrmot/model.py
+++ b/src/ewatercycle/plugins/marrmot/model.py
@@ -4,7 +4,7 @@
 import logging
 from dataclasses import asdict, dataclass
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, ItemsView, Literal
 
 import scipy.io as sio
 from pydantic import PrivateAttr, model_validator
@@ -136,7 +136,7 @@ def _make_cfg_file(self, **kwargs) -> Path:
         return config_file
 
     @property
-    def parameters(self) -> dict[str, Any]:
+    def parameters(self) -> ItemsView[str, Any]:
         """List MarrmotM01's parameters and their values.
 
         Model parameters:
@@ -157,7 +157,7 @@ def parameters(self) -> dict[str, Any]:
             "solver": get_solver(self._config),
             "start time": get_marrmot_time(self._config, "start"),
             "end time": get_marrmot_time(self._config, "end"),
-        }
+        }.items()
 
 
 M14_PARAMS = (
@@ -272,7 +272,7 @@ def _make_cfg_file(self, **kwargs) -> Path:
         return config_file
 
     @property
-    def parameters(self) -> dict[str, Any]:
+    def parameters(self) -> ItemsView[str, Any]:
         """List the parameters for this model.
 
         Exposed Marrmot M14 parameters:
@@ -304,4 +304,4 @@ def parameters(self) -> dict[str, Any]:
                 "end time": get_marrmot_time(self._config, "end"),
             }
         )
-        return pars
+        return pars.items()
diff --git a/src/ewatercycle/plugins/pcrglobwb/model.py b/src/ewatercycle/plugins/pcrglobwb/model.py
index ecde6b02..edee4293 100644
--- a/src/ewatercycle/plugins/pcrglobwb/model.py
+++ b/src/ewatercycle/plugins/pcrglobwb/model.py
@@ -2,7 +2,7 @@
 
 import logging
 from os import PathLike
-from typing import Any, Optional
+from typing import Any, ItemsView, Optional
 
 import bmipy
 import numpy as np
@@ -97,7 +97,7 @@ def _initialize_config(self: "PCRGlobWB") -> "PCRGlobWB":
         return self
 
     @property
-    def parameters(self) -> dict[str, Any]:
+    def parameters(self) -> ItemsView[str, Any]:
         """List the parameters for this model."""
         return {
             "start_time": f"{self._config.get('globalOptions', 'startTime')}T00:00:00Z",
@@ -106,7 +106,7 @@ def parameters(self) -> dict[str, Any]:
             "max_spinups_in_years": self._config.get(
                 "globalOptions", "maxSpinUpsInYears"
             ),
-        }
+        }.items()
 
     def _make_cfg_file(self, **kwargs):
         self._update_config(**kwargs)
diff --git a/src/ewatercycle/plugins/wflow/model.py b/src/ewatercycle/plugins/wflow/model.py
index b71b2922..c12a8469 100644
--- a/src/ewatercycle/plugins/wflow/model.py
+++ b/src/ewatercycle/plugins/wflow/model.py
@@ -4,7 +4,7 @@
 import logging
 import shutil
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, ItemsView, Optional
 
 import bmipy
 import numpy as np
@@ -142,11 +142,11 @@ def _make_cfg_dir(self, cfg_dir: Optional[str] = None, **kwargs) -> Path:
         return self._work_dir
 
     @property
-    def parameters(self) -> dict[str, Any]:
+    def parameters(self) -> ItemsView[str, Any]:
         return {
             "start_time": _wflow_to_iso(self._config.get("run", "starttime")),
             "end_time": _wflow_to_iso(self._config.get("run", "endtime")),
-        }
+        }.items()
 
 
 def _wflow_to_iso(time):
diff --git a/src/ewatercycle/testing/fake_models.py b/src/ewatercycle/testing/fake_models.py
index 002e223f..044b6d1e 100644
--- a/src/ewatercycle/testing/fake_models.py
+++ b/src/ewatercycle/testing/fake_models.py
@@ -1,4 +1,6 @@
 """Fake BMI models."""
+from unittest.mock import Mock
+
 import numpy as np
 from bmipy import Bmi
 
@@ -135,3 +137,135 @@ def get_grid_nodes_per_face(
 
     def get_grid_face_edges(self, grid: int, face_edges: np.ndarray) -> np.ndarray:
         raise self.exc
+
+
+class NotImplementedModel(FailingModel):
+    def __init__(self, exc=NotImplementedError()):
+        super().__init__(exc=exc)
+
+
+class WithMocksMixin:
+    """Mock the bmi methods that return None and have no getter companion.
+
+    Use `instance.mock.<method name>.assert_called_once_with()`
+    to check if the method is called correctly.
+    """
+
+    def __init__(self):
+        self.mock = Mock()
+
+    def initialize(self, config_file: str) -> None:
+        self.mock.initialize(config_file)
+
+    def finalize(self):
+        self.mock.finalize()
+
+
+class WithDailyMixin:
+    """Mock the bmi methods that deal wtih time.
+
+    Behaves like a daily model which started since epoch.
+    """
+
+    def __init__(self) -> None:
+        self.time = 0.0
+
+    def update(self):
+        self.time = self.time + self.get_time_step()
+
+    def get_current_time(self):
+        return self.time
+
+    def get_start_time(self):
+        return 0.0
+
+    def get_end_time(self):
+        return 100.0
+
+    def get_time_step(self):
+        return 1.0
+
+    def get_time_units(self):
+        return "days since 1970-01-01"
+
+
+class DummyModelWith2DRectilinearGrid(
+    WithMocksMixin, WithDailyMixin, NotImplementedModel
+):
+    def __init__(self):
+        super().__init__()
+        # not sure why extra call to init is needed,
+        # but without the self.time is not initialized
+        WithMocksMixin.__init__(self)
+        WithDailyMixin.__init__(self)
+        self.dtype = np.dtype("float32")
+        self.value = np.array(
+            [
+                1.1,
+                2.2,
+                3.3,
+                4.4,
+                5.5,
+                6.6,
+                7.7,
+                8.8,
+                9.9,
+                10.1,
+                11.1,
+                12.1,
+            ],
+            dtype=self.dtype,
+        )
+
+    def get_output_var_names(self) -> tuple[str]:
+        return ("plate_surface__temperature",)
+
+    def get_var_type(self, name):
+        return str(self.dtype)
+
+    def get_var_grid(self, name):
+        return 0
+
+    def get_var_units(self, name):
+        return "K"
+
+    def get_var_itemsize(self, name):
+        return self.dtype.itemsize
+
+    def get_var_nbytes(self, name):
+        return self.dtype.itemsize * self.value.size
+
+    def get_grid_type(self, grid):
+        return "rectilinear"
+
+    def get_grid_size(self, grid):
+        return 12  # 4 longs * 3 lats
+
+    def get_grid_rank(self, grid: int) -> int:
+        return 2
+
+    def get_grid_shape(self, grid: int, shape: np.ndarray) -> np.ndarray:
+        np.copyto(src=[3, 4], dst=shape)
+        return shape
+
+    def get_value(self, name, dest):
+        np.copyto(src=self.value, dst=dest)
+        return dest
+
+    def get_value_at_indices(self, name, dest, inds):
+        np.copyto(src=self.value[inds], dst=dest)
+        return dest
+
+    def set_value(self, name, src):
+        self.value[:] = src
+
+    def set_value_at_indices(self, name, inds, src):
+        self.value[inds] = src
+
+    def get_grid_x(self, grid: int, x: np.ndarray) -> np.ndarray:
+        np.copyto(src=[0.1, 0.2, 0.3, 0.4], dst=x)
+        return x
+
+    def get_grid_y(self, grid: int, y: np.ndarray) -> np.ndarray:
+        np.copyto(src=[1.1, 1.2, 1.3], dst=y)
+        return y
diff --git a/tests/plugins/hype/test_model.py b/tests/plugins/hype/test_model.py
index 577f8e38..675ac6a8 100644
--- a/tests/plugins/hype/test_model.py
+++ b/tests/plugins/hype/test_model.py
@@ -6,6 +6,7 @@
 import numpy as np
 import pytest
 from grpc4bmi.bmi_client_apptainer import BmiClientApptainer
+from grpc4bmi.bmi_optionaldest import OptionalDestBmi
 
 from ewatercycle import CFG
 from ewatercycle.base.parameter_set import ParameterSet
@@ -140,7 +141,7 @@ def test_parameters(self, model):
             "start_time": "1961-01-01T00:00:00Z",
             "end_time": "1963-12-31T00:00:00Z",
             "crit_time": "1962-01-01T00:00:00Z",
-        }
+        }.items()
         assert model.parameters == expected
 
     def test_get_value_as_xarray(self, model):
@@ -154,12 +155,12 @@ class MockedBmi(FailingModel):
             def get_var_grid(self, name):
                 return 1
 
-            def get_grid_x(self, grid_id):
+            def get_grid_x(self, grid_id, dest):
                 return np.array(
                     [5.8953929, 4.9553967, 5.6387277]
                 )  # x subbasin lons of subbasinsin hype
 
-            def get_grid_y(self, grid_id):
+            def get_grid_y(self, grid_id, dest):
                 return np.array(
                     [51.16437912, 50.21104813, 48.6910553]
                 )  # y are lats of subbasins in hype
@@ -177,11 +178,20 @@ def get_var_itemsize(self, name):
             def get_var_nbytes(self, name):
                 return np.float64().size * 3 * 3
 
-        model._bmi = MockedBmi()
+            def get_grid_rank(self, grid):
+                return 2
+
+            def get_grid_shape(self, grid, shape):
+                return (3, 3)
+
+            def get_grid_type(self, grid):
+                return "rectilinear"
+
+        model._bmi = OptionalDestBmi(MockedBmi())
 
         actual = model.get_value_at_coords("comp outflow olake", lon=[5], lat=[50])
         assert actual == np.array([13.0])
-        assert model._bmi.indices == [1]
+        assert model.bmi.origin.indices == [1]
 
 
 class TestWithOnlyParameterSetAndFullSetup:
@@ -250,7 +260,7 @@ def test_parameters(self, model_with_setup, model):
             "start_time": "2000-01-01T00:00:00Z",
             "end_time": "2010-12-31T00:00:00Z",
             "crit_time": "2002-01-01T00:00:00Z",
-        }
+        }.items()
         assert model.parameters == expected
 
 
@@ -393,5 +403,5 @@ def test_parameters(self, model):
             "start_time": "1986-01-02T00:00:00Z",
             "end_time": "2018-01-02T00:00:00Z",
             "crit_time": "1986-01-02T00:00:00Z",
-        }
+        }.items()
         assert model.parameters == expected
diff --git a/tests/plugins/lisflood/test_model.py b/tests/plugins/lisflood/test_model.py
index bcf271b7..55556e25 100644
--- a/tests/plugins/lisflood/test_model.py
+++ b/tests/plugins/lisflood/test_model.py
@@ -7,6 +7,7 @@
 import numpy as np
 import pytest
 from grpc4bmi.bmi_client_apptainer import BmiClientApptainer
+from grpc4bmi.bmi_optionaldest import OptionalDestBmi
 from numpy.testing import assert_array_equal
 from pytest import TempPathFactory
 
@@ -84,7 +85,7 @@ def test_default_parameters(self, model: Lisflood):
             "MaskMap": "$(PathMaps)/masksmall.map",
             "start_time": "1986-01-02T00:00:00Z",
             "end_time": "2018-01-02T00:00:00Z",
-        }
+        }.items()
         assert model.parameters == expected_parameters
 
     @pytest.fixture
@@ -140,7 +141,7 @@ def test_setup(self, model_with_setup, tmp_path):
 
     class TestGetValueAtCoords:
         def test_get_value_at_coords_single(self, caplog, model: Lisflood):
-            model._bmi = MockedBmi()
+            model._bmi = OptionalDestBmi(MockedBmi())
 
             with caplog.at_level(logging.DEBUG):
                 result = model.get_value_at_coords(
@@ -155,19 +156,19 @@ def test_get_value_at_coords_single(self, caplog, model: Lisflood):
             assert msg[0] in caplog.text
             assert msg[1] in caplog.text
             assert result == np.array([1.0])
-            assert model._bmi.indices == [311]
+            assert model.bmi.origin.indices == [311]
 
         def test_get_value_at_coords_multiple(self, model: Lisflood):
-            model._bmi = MockedBmi()
+            model._bmi = OptionalDestBmi(MockedBmi())
             model.get_value_at_coords(
                 "Discharge",
                 lon=[-124.45, -124.35, -121.45],
                 lat=[53.95, 52.93, 52.65],
             )
-            assert_array_equal(model._bmi.indices, [0, 311, 433])
+            assert_array_equal(model.bmi.origin.indices, [0, 311, 433])
 
         def test_get_value_at_coords_faraway(self, model: Lisflood):
-            model._bmi = MockedBmi()
+            model._bmi = OptionalDestBmi(MockedBmi())
             with pytest.raises(ValueError) as excinfo:
                 model.get_value_at_coords("Discharge", lon=[0.0], lat=[0.0])
             msg = str(excinfo.value)
@@ -248,7 +249,7 @@ def test_parameters_after_setup(self, model_with_setup, tmp_path):
                 "MaskMap": f"{tmp_path}/custommask/mask",
                 "start_time": "1986-01-02T00:00:00Z",
                 "end_time": "2018-01-02T00:00:00Z",
-            }
+            }.items()
             assert model_with_setup[3].parameters == expected_parameters
 
 
@@ -258,10 +259,10 @@ class MockedBmi(FailingModel):
     def get_var_grid(self, name):
         return 0
 
-    def get_grid_shape(self, grid_id):
+    def get_grid_shape(self, grid_id, dest):
         return 14, 31  # shape returns (len(y), len(x))
 
-    def get_grid_x(self, grid_id):
+    def get_grid_x(self, grid_id, dest):
         return np.array(
             [
                 -124.450000000003,
@@ -298,7 +299,7 @@ def get_grid_x(self, grid_id):
             ]
         )
 
-    def get_grid_y(self, grid_id):
+    def get_grid_y(self, grid_id, dest):
         return np.array(
             [
                 53.950000000002,
@@ -333,3 +334,9 @@ def get_var_itemsize(self, name):
 
     def get_var_nbytes(self, name):
         return np.float64().size * 14 * 31
+
+    def get_grid_rank(self, grid):
+        return 2
+
+    def get_grid_type(self, grid):
+        return "rectilinear"
diff --git a/tests/plugins/marrmot/test_model_m01.py b/tests/plugins/marrmot/test_model_m01.py
index d0ca9fbf..45c05b06 100644
--- a/tests/plugins/marrmot/test_model_m01.py
+++ b/tests/plugins/marrmot/test_model_m01.py
@@ -70,7 +70,7 @@ def test_parameters(self, model):
             "solver": Solver(),
             "start time": "1989-01-01T00:00:00Z",
             "end time": "1992-12-31T00:00:00Z",
-        }
+        }.items()
         assert model.parameters == expected
 
     def test_setup(self, model_with_setup, forcing_file):
@@ -98,7 +98,7 @@ def test_parameters_after_setup(self, model_with_setup):
             "solver": Solver(),
             "start time": "1989-01-01T00:00:00Z",
             "end time": "1992-12-31T00:00:00Z",
-        }
+        }.items()
         assert model.parameters == expected
 
     def test_get_value_as_xarray(self, model_with_setup):
diff --git a/tests/plugins/marrmot/test_model_m14.py b/tests/plugins/marrmot/test_model_m14.py
index 2939e606..a9fd07a1 100644
--- a/tests/plugins/marrmot/test_model_m14.py
+++ b/tests/plugins/marrmot/test_model_m14.py
@@ -61,7 +61,7 @@ def test_parameters(self, model):
             "solver": Solver(),
             "start time": "1989-01-01T00:00:00Z",
             "end time": "1992-12-31T00:00:00Z",
-        }
+        }.items()
         assert model.parameters == expected
 
     def test_setup(self, model_with_setup):
@@ -99,7 +99,7 @@ def test_parameters_after_setup(self, model_with_setup):
             "solver": Solver(),
             "start time": "1989-01-01T00:00:00Z",
             "end time": "1992-12-31T00:00:00Z",
-        }
+        }.items()
         assert model.parameters == expected
 
     def test_get_value_as_xarray(self, model_with_setup):
diff --git a/tests/plugins/wflow/test_model.py b/tests/plugins/wflow/test_model.py
index 0452ffa7..657b21f5 100644
--- a/tests/plugins/wflow/test_model.py
+++ b/tests/plugins/wflow/test_model.py
@@ -8,6 +8,7 @@
 import pytest
 from grpc import FutureTimeoutError
 from grpc4bmi.bmi_client_apptainer import BmiClientApptainer
+from grpc4bmi.bmi_optionaldest import OptionalDestBmi
 
 from ewatercycle import CFG
 from ewatercycle.base.parameter_set import ParameterSet
@@ -28,13 +29,13 @@ def initialize(self, config_file):
     def get_var_grid(self, name):
         return 1
 
-    def get_grid_shape(self, grid_id):
+    def get_grid_shape(self, grid_id, dest):
         return 3, 2
 
-    def get_grid_y(self, grid_id):
+    def get_grid_y(self, grid_id, dest):
         return np.array([45.0, 46.0, 47.0])
 
-    def get_grid_x(self, grid_id):
+    def get_grid_x(self, grid_id, dest):
         return np.array([5.0, 6.0])
 
     def get_grid_spacing(self, grid_id):
@@ -53,6 +54,12 @@ def get_var_itemsize(self, name):
     def get_var_nbytes(self, name):
         return np.float64().size * 3 * 2
 
+    def get_grid_rank(self, grid):
+        return 2
+
+    def get_grid_type(self, grid):
+        return "rectilinear"
+
 
 @pytest.fixture
 def mocked_config(tmp_path: Path):
@@ -106,7 +113,7 @@ def model(parameter_set):
 @pytest.fixture
 def initialized_model(model):
     """Model with fake parameterset and fake BMI instance."""
-    model._bmi = MockedBmi()
+    model._bmi = OptionalDestBmi(MockedBmi())
     return model
 
 
@@ -196,4 +203,4 @@ def test_get_value_as_coords(initialized_model, caplog):
 
     assert msg1 in caplog.text and msg2 in caplog.text
     assert result == np.array([1.0])
-    assert model._bmi.indices == [4]
+    assert model.bmi.origin.indices == [4]
diff --git a/tests/src/base/__init__.py b/tests/src/base/__init__.py
new file mode 100644
index 00000000..490920ff
--- /dev/null
+++ b/tests/src/base/__init__.py
@@ -0,0 +1 @@
+"""Tests for ewatercycle.base module."""
diff --git a/tests/src/base/test_model.py b/tests/src/base/test_model.py
new file mode 100644
index 00000000..862ab9b8
--- /dev/null
+++ b/tests/src/base/test_model.py
@@ -0,0 +1,397 @@
+from datetime import datetime
+from pathlib import Path
+from typing import Type
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+import xarray as xr
+from bmipy import Bmi
+from grpc4bmi.bmi_optionaldest import OptionalDestBmi
+from numpy.testing import assert_almost_equal
+from pydantic import ConfigDict
+from pytest import fixture
+from xarray.testing import assert_allclose
+
+from ewatercycle.base.forcing import DefaultForcing
+from ewatercycle.base.model import ContainerizedModel, LocalModel, eWaterCycleModel
+from ewatercycle.base.parameter_set import ParameterSet
+from ewatercycle.testing.fake_models import DummyModelWith2DRectilinearGrid
+
+
+class MockModel(eWaterCycleModel):
+    # Instead of using custom constructor had to do it the Pydantic way
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    mybmi: Bmi
+
+    def _make_bmi_instance(self):
+        return OptionalDestBmi(self.mybmi)
+
+
+@fixture
+def mocked_bmi():
+    return DummyModelWith2DRectilinearGrid()
+
+
+@fixture
+def mocked_model(mocked_config, mocked_bmi):
+    return MockModel(mybmi=mocked_bmi)
+
+
+def test_version(mocked_model: eWaterCycleModel):
+    assert mocked_model.version == ""
+
+
+def test_parameters(mocked_model: eWaterCycleModel):
+    assert mocked_model.parameters == {}.items()
+
+
+class TestWithSetup:
+    @fixture(autouse=True)
+    def setup_on_mocked_model(self, mocked_model: eWaterCycleModel, tmp_path: Path):
+        return mocked_model.setup(cfg_dir=str(tmp_path))
+
+    def test_setup_cfg_dir(self, setup_on_mocked_model, tmp_path: Path):
+        assert setup_on_mocked_model[1] == str(tmp_path)
+
+    def test_setup_cfg_file(self, setup_on_mocked_model, tmp_path: Path):
+        expected = str(tmp_path / "config.yaml")
+        assert setup_on_mocked_model[0] == expected
+
+    def test_setup_cfg_file_content(self, setup_on_mocked_model, tmp_path: Path):
+        expected = "{}"
+        assert (tmp_path / "config.yaml").read_text().strip() == expected
+
+    def test_initialize(self, mocked_model: eWaterCycleModel, setup_on_mocked_model):
+        config_file = setup_on_mocked_model[0]
+
+        mocked_model.initialize(config_file)
+
+        mocked_model.bmi.origin.mock.initialize.assert_called_once_with(config_file)
+
+    def test_finalize(
+        self,
+        mocked_model: eWaterCycleModel,
+        mocked_bmi: DummyModelWith2DRectilinearGrid,
+    ):
+        assert hasattr(mocked_model, "_bmi")
+
+        mocked_model.finalize()
+
+        with pytest.raises(AttributeError, match="has no attribute '_bmi'"):
+            mocked_model.bmi
+        mocked_bmi.mock.finalize.assert_called_once_with()
+
+    def test_update(self, mocked_model: eWaterCycleModel):
+        mocked_model.update()
+        mocked_model.update()
+        mocked_model.update()
+
+        assert mocked_model.time == pytest.approx(3.0)
+
+    def test_time(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.time
+
+        assert result == pytest.approx(0.0)
+
+    def test_time_step(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.time_step
+
+        assert result == pytest.approx(1.0)
+
+    def test_start_time(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.start_time
+
+        assert result == pytest.approx(0.0)
+
+    def test_end_time(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.end_time
+
+        assert result == pytest.approx(100.0)
+
+    def test_time_as_datetime(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.time_as_datetime
+
+        assert result == datetime(1970, 1, 1)
+
+    def test_start_time_as_datetime(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.start_time_as_datetime
+
+        assert result == datetime(1970, 1, 1)
+
+    def test_end_time_as_datetime(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.end_time_as_datetime
+
+        # 100th day of 1970
+        assert result == datetime(1970, 4, 11)
+
+    def test_time_as_isostr(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.time_as_isostr
+
+        assert result == "1970-01-01T00:00:00Z"
+
+    def test_start_time_as_isostr(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.start_time_as_isostr
+
+        assert result == "1970-01-01T00:00:00Z"
+
+    def test_end_time_as_isostr(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.end_time_as_isostr
+
+        assert result == "1970-04-11T00:00:00Z"
+
+    def test_time_units(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.time_units
+
+        assert result == "days since 1970-01-01"
+
+    def test_output_var_names(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.output_var_names
+
+        assert result == ("plate_surface__temperature",)
+
+    def test_get_value(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.get_value("plate_surface__temperature")
+
+        expected = np.array(
+            [
+                1.1,
+                2.2,
+                3.3,
+                4.4,
+                5.5,
+                6.6,
+                7.7,
+                8.8,
+                9.9,
+                10.1,
+                11.1,
+                12.1,
+            ],
+            dtype=np.float32,
+        )
+        assert_almost_equal(result, expected)
+
+    def test_get_value_at_coords(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.get_value_at_coords(
+            "plate_surface__temperature",
+            lat=[
+                1.2,
+            ],
+            lon=[
+                0.3,
+            ],
+        )
+
+        expected = np.array(
+            [
+                7.7,
+            ],
+            dtype=np.float32,
+        )
+        assert_almost_equal(result, expected)
+
+    def test_set_value(
+        self,
+        mocked_model: eWaterCycleModel,
+        mocked_bmi: DummyModelWith2DRectilinearGrid,
+    ):
+        new_values = np.full((12,), 4.2, dtype=np.float32)
+
+        mocked_model.set_value("plate_surface__temperature", new_values)
+
+        assert_almost_equal(mocked_bmi.value, new_values)
+
+    def test_set_value_at_coords(
+        self,
+        mocked_model: eWaterCycleModel,
+        mocked_bmi: DummyModelWith2DRectilinearGrid,
+    ):
+        new_values = np.array([4.2, 1.23], dtype=np.float32)
+
+        mocked_model.set_value_at_coords(
+            "plate_surface__temperature",
+            lon=[0.3, 0.2],
+            lat=[1.2, 1.1],
+            values=new_values,
+        )
+
+        expected = np.array(
+            [
+                1.1,
+                1.23,
+                3.3,
+                4.4,
+                5.5,
+                6.6,
+                4.2,
+                8.8,
+                9.9,
+                10.1,
+                11.1,
+                12.1,
+            ],
+            dtype=np.float32,
+        )
+        assert_almost_equal(mocked_bmi.value, expected)
+
+    def test_get_value_as_xarray(self, mocked_model: eWaterCycleModel):
+        result = mocked_model.get_value_as_xarray("plate_surface__temperature")
+
+        expected = xr.DataArray(
+            data=np.array(
+                [
+                    [
+                        [1.1, 2.2, 3.3, 4.4],
+                        [5.5, 6.6, 7.7, 8.8],
+                        [9.9, 10.1, 11.1, 12.1],
+                    ]
+                ],
+                dtype=np.float32,
+            ),
+            coords={
+                "longitude": [0.1, 0.2, 0.3, 0.4],
+                "latitude": [1.1, 1.2, 1.3],
+                "time": [datetime(1970, 1, 1)],
+            },
+            dims=["time", "latitude", "longitude"],
+            name="plate_surface__temperature",
+            attrs={"units": "K"},
+        )
+        assert_allclose(result, expected)
+        # Check that data and coords are aligned
+        assert result.sel(
+            longitude=0.2, latitude=1.2, time=datetime(1970, 1, 1)
+        ).values.tolist() == pytest.approx(6.6)
+
+
+class DummyLocalModel(LocalModel):
+    bmi_class: Type[Bmi] = DummyModelWith2DRectilinearGrid
+
+
+# bit ugly to have version here,
+# but that is how version of LocalModel class is determined
+__version__ = "1.2.3"
+
+
+class TestLocalModel:
+    @fixture
+    def model(self):
+        return DummyLocalModel()
+
+    def test_version(self, model):
+        assert model.version == "1.2.3"
+
+    @fixture(autouse=True)
+    def setup_on_mocked_model(self, model: eWaterCycleModel, tmp_path: Path):
+        return model.setup(cfg_dir=str(tmp_path))
+
+    def test_bmi_class(self, model: eWaterCycleModel):
+        assert isinstance(model.bmi, OptionalDestBmi)
+        assert isinstance(model.bmi.origin, DummyModelWith2DRectilinearGrid)
+
+
+class TestContainerizedModel:
+    def test_version(self):
+        model = ContainerizedModel(bmi_image="ewatercycle/ewatercycle_dummy:latest")
+
+        assert model.version == "latest"
+
+    @patch("ewatercycle.base.model.start_container")
+    def test_setup(self, mocked_start_container, tmp_path: Path):
+        model = ContainerizedModel(bmi_image="ewatercycle/ewatercycle_dummy:latest")
+
+        model.setup(cfg_dir=str(tmp_path))
+
+        mocked_start_container.assert_called_once_with(
+            image="ewatercycle/ewatercycle_dummy:latest",
+            work_dir=tmp_path,
+            input_dirs=[],
+            timeout=300,
+        )
+
+    @patch("ewatercycle.base.model.start_container")
+    def test_setup_with_additional_input_dirs(
+        self, mocked_start_container, tmp_path: Path
+    ):
+        forcing_dir = tmp_path / "forcing"
+        forcing = DefaultForcing(
+            directory=forcing_dir,
+            start_time="1989-01-02T00:00:00Z",
+            end_time="1999-01-02T00:00:00Z",
+        )
+        parameter_set_dir = tmp_path / "parameter_set"
+        parameter_set = ParameterSet(
+            name="test",
+            directory=parameter_set_dir,
+            config="config.yaml",
+            target_model="containerizedmodel",
+            supported_model_versions=["latest"],
+        )
+        model = ContainerizedModel(
+            bmi_image="ewatercycle/ewatercycle_dummy:latest",
+            parameter_set=parameter_set,
+            forcing=forcing,
+        )
+
+        model.setup(cfg_dir=str(tmp_path))
+
+        mocked_start_container.assert_called_once_with(
+            image="ewatercycle/ewatercycle_dummy:latest",
+            work_dir=tmp_path,
+            input_dirs=[
+                str(parameter_set_dir),
+                str(forcing_dir),
+            ],
+            timeout=300,
+        )
+
+
+class VersionedMockModel(MockModel):
+    @property
+    def version(self) -> str:
+        return "latest"
+
+
+class TestParameterSetValidation:
+    def test_with_matching_model_and_version(self, mocked_config, mocked_bmi):
+        parameter_set = ParameterSet(
+            name="test",
+            directory=".",
+            config="config.yaml",
+            target_model="versionedmockmodel",
+            supported_model_versions=["latest"],
+        )
+
+        model = VersionedMockModel(mybmi=mocked_bmi, parameter_set=parameter_set)
+        assert (
+            repr(model)
+            == "VersionedMockModel(parameter_set=ParameterSet(name='test', directory=PosixPath('.'), config=PosixPath('config.yaml'), doi='N/A', target_model='versionedmockmodel', supported_model_versions={'latest'}, downloader=None), forcing=None)"
+        )
+
+    def test_with_mismatching_model(self, mocked_config, mocked_bmi):
+        parameter_set = ParameterSet(
+            name="test",
+            directory=".",
+            config="config.yaml",
+            target_model="wrongmodel",
+            supported_model_versions=["latest"],
+        )
+
+        with pytest.raises(ValueError, match="Parameter set has wrong target model"):
+            VersionedMockModel(mybmi=mocked_bmi, parameter_set=parameter_set)
+
+    def test_with_mismatching_version(self, mocked_config, mocked_bmi):
+        parameter_set = ParameterSet(
+            name="test",
+            directory=".",
+            config="config.yaml",
+            target_model="versionedmockmodel",
+            supported_model_versions=["1.0.0"],
+        )
+
+        with pytest.raises(
+            ValueError,
+            match="Parameter set 'test' not compatible with this model version.",
+        ):
+            VersionedMockModel(mybmi=mocked_bmi, parameter_set=parameter_set)
diff --git a/tests/src/fake_models.py b/tests/src/fake_models.py
deleted file mode 100644
index c6fc8304..00000000
--- a/tests/src/fake_models.py
+++ /dev/null
@@ -1,432 +0,0 @@
-# TODO do not copy from https://github.com/eWaterCycle/grpc4bmi/blob/main/test/fake_models.py
-# but import instead
-from typing import Tuple
-
-import numpy
-import numpy as np
-from bmipy import Bmi
-from grpc4bmi.constants import GRPC_MAX_MESSAGE_LENGTH
-
-
-class SomeException(Exception):
-    pass
-
-
-class FailingModel(Bmi):
-    def __init__(self, exc):
-        self.exc = exc
-
-    def initialize(self, filename):
-        raise self.exc
-
-    def update(self):
-        raise self.exc
-
-    def update_until(self, time: float) -> None:
-        raise self.exc
-
-    def finalize(self):
-        raise self.exc
-
-    def get_component_name(self):
-        raise self.exc
-
-    def get_input_item_count(self) -> int:
-        raise self.exc
-
-    def get_output_item_count(self) -> int:
-        raise self.exc
-
-    def get_input_var_names(self):
-        raise self.exc
-
-    def get_output_var_names(self):
-        raise self.exc
-
-    def get_start_time(self):
-        raise self.exc
-
-    def get_current_time(self):
-        raise self.exc
-
-    def get_end_time(self):
-        raise self.exc
-
-    def get_time_step(self):
-        raise self.exc
-
-    def get_time_units(self):
-        raise self.exc
-
-    def get_var_type(self, name):
-        raise self.exc
-
-    def get_var_units(self, name):
-        raise self.exc
-
-    def get_var_itemsize(self, name):
-        raise self.exc
-
-    def get_var_nbytes(self, name):
-        raise self.exc
-
-    def get_var_grid(self, name):
-        raise self.exc
-
-    def get_value(self, name, dest):
-        raise self.exc
-
-    def get_value_ptr(self, name):
-        raise self.exc
-
-    def get_value_at_indices(self, name, dest, inds):
-        raise self.exc
-
-    def set_value(self, name, src):
-        raise self.exc
-
-    def set_value_at_indices(self, name, inds, src):
-        raise self.exc
-
-    def get_grid_shape(self, grid, shape):
-        raise self.exc
-
-    def get_grid_x(self, grid, x):
-        raise self.exc
-
-    def get_grid_y(self, grid, y):
-        raise self.exc
-
-    def get_grid_z(self, grid, z):
-        raise self.exc
-
-    def get_grid_spacing(self, grid, spacing):
-        raise self.exc
-
-    def get_grid_origin(self, grid, origin):
-        raise self.exc
-
-    def get_grid_rank(self, grid):
-        raise self.exc
-
-    def get_grid_size(self, grid):
-        raise self.exc
-
-    def get_grid_type(self, grid):
-        raise self.exc
-
-    def get_var_location(self, name: str) -> str:
-        raise self.exc
-
-    def get_grid_node_count(self, grid: int) -> int:
-        raise self.exc
-
-    def get_grid_edge_count(self, grid: int) -> int:
-        raise self.exc
-
-    def get_grid_face_count(self, grid: int) -> int:
-        raise self.exc
-
-    def get_grid_edge_nodes(self, grid: int, edge_nodes: np.ndarray) -> np.ndarray:
-        raise self.exc
-
-    def get_grid_face_nodes(self, grid: int, face_nodes: np.ndarray) -> np.ndarray:
-        raise self.exc
-
-    def get_grid_nodes_per_face(
-        self, grid: int, nodes_per_face: np.ndarray
-    ) -> np.ndarray:
-        raise self.exc
-
-    def get_grid_face_edges(self, grid: int, face_edges: np.ndarray) -> np.ndarray:
-        raise self.exc
-
-
-class GridModel(FailingModel):
-    def __init__(self):
-        super(GridModel, self).__init__(SomeException("not used"))
-
-    def initialize(self, filename):
-        pass
-
-    def get_output_var_names(self) -> Tuple[str]:
-        return ("plate_surface__temperature",)
-
-    def get_var_grid(self, name):
-        return 0
-
-
-class UniRectGridModel(GridModel):
-    def get_grid_type(self, grid):
-        return "uniform_rectilinear"
-
-    def get_grid_rank(self, grid):
-        return 3
-
-    def get_grid_size(self, grid):
-        return 24
-
-    def get_grid_shape(self, grid: int, shape: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[2, 3, 4], dst=shape)
-        return shape
-
-    def get_grid_origin(self, grid, dest):
-        numpy.copyto(src=[0.1, 1.1, 2.1], dst=dest)
-        return dest
-
-    def get_grid_spacing(self, grid, dest):
-        numpy.copyto(src=[0.1, 0.2, 0.3], dst=dest)
-        return dest
-
-
-class Rect3DGridModel(GridModel):
-    def get_grid_type(self, grid):
-        return "rectilinear"
-
-    def get_grid_size(self, grid):
-        return 24
-
-    def get_grid_rank(self, grid: int) -> int:
-        return 3
-
-    def get_grid_shape(self, grid: int, shape: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[2, 3, 4], dst=shape)
-        return shape
-
-    def get_grid_x(self, grid: int, x: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[0.1, 0.2, 0.3, 0.4], dst=x)
-        return x
-
-    def get_grid_y(self, grid: int, y: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[1.1, 1.2, 1.3], dst=y)
-        return y
-
-    def get_grid_z(self, grid: int, z: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[2.1, 2.2], dst=z)
-        return z
-
-
-class Rect2DGridModel(Rect3DGridModel):
-    def get_grid_size(self, grid):
-        return 12  # 4*3
-
-    def get_grid_rank(self, grid: int) -> int:
-        return 2
-
-    def get_grid_shape(self, grid: int, shape: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[3, 4], dst=shape)
-        return shape
-
-    def get_grid_z(self, grid: int, z: np.ndarray) -> np.ndarray:
-        raise NotImplementedError("Do not know what z is")
-
-
-class Structured3DQuadrilateralsGridModel(GridModel):
-    # Grid shape:
-    #    0
-    #   / \
-    #  /   \
-    # 3     1
-    #  \    /
-    #   \  /
-    #    2
-    #
-    def get_grid_type(self, grid):
-        return "structured_quadrilateral"
-
-    def get_grid_rank(self, grid: int) -> int:
-        return 3
-
-    def get_grid_size(self, grid):
-        return 4
-
-    def get_grid_shape(self, grid, shape):
-        numpy.copyto(src=[1, 2, 2], dst=shape)
-        return shape
-
-    def get_grid_x(self, grid, x):
-        numpy.copyto(src=[1.1, 0.1, 1.1, 2.1], dst=x)
-        return x
-
-    def get_grid_y(self, grid: int, y: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[2.2, 1.2, 0.2, 2.2], dst=y)
-        return y
-
-    def get_grid_z(self, grid: int, z: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[1.1, 2.2, 3.3, 4.4], dst=z)
-        return z
-
-
-class Structured2DQuadrilateralsGridModel(GridModel):
-    # Grid shape:
-    #    0
-    #   / \
-    #  /   \
-    # 3     1
-    #  \    /
-    #   \  /
-    #    2
-    #
-    def get_grid_type(self, grid):
-        return "structured_quadrilateral"
-
-    def get_grid_rank(self, grid: int) -> int:
-        return 2
-
-    def get_grid_size(self, grid):
-        return 4
-
-    def get_grid_shape(self, grid, shape):
-        numpy.copyto(src=[2, 2], dst=shape)
-        return shape
-
-    def get_grid_x(self, grid, x):
-        numpy.copyto(src=[1.1, 0.1, 1.1, 2.1], dst=x)
-        return x
-
-    def get_grid_y(self, grid: int, y: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[2.2, 1.2, 0.2, 2.2], dst=y)
-        return y
-
-    def get_grid_z(self, grid: int, z: np.ndarray) -> np.ndarray:
-        raise NotImplementedError("Do not know what z is")
-
-
-class UnstructuredGridBmiModel(GridModel):
-    # Uses grid example at https://bmi.readthedocs.io/en/latest/model_grids.html#unstructured-grids
-    # Grid shape:
-    #    3-----\
-    #   / \     \
-    # 0    \    \---5
-    # \     2--/   /
-    #  \    /     /
-    #   \  /     /
-    #    1---\  /
-    #         4
-    def get_grid_type(self, grid):
-        return "unstructured"
-
-    def get_grid_shape(self, grid, dest):
-        raise NotImplementedError("Do not know what shape is")
-
-    def get_grid_size(self, grid):
-        return 6
-
-    def get_grid_rank(self, grid: int) -> int:
-        return 2
-
-    def get_grid_node_count(self, grid: int) -> int:
-        return 6
-
-    def get_grid_edge_count(self, grid: int) -> int:
-        return 8
-
-    def get_grid_face_count(self, grid: int) -> int:
-        return 3
-
-    def get_grid_edge_nodes(self, grid: int, edge_nodes: np.ndarray) -> np.ndarray:
-        numpy.copyto(
-            src=(0, 1, 1, 2, 2, 3, 3, 0, 1, 4, 4, 5, 5, 2, 5, 3), dst=edge_nodes
-        )
-        return edge_nodes
-
-    def get_grid_face_nodes(self, grid: int, face_nodes: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=(0, 1, 2, 3, 1, 4, 5, 2, 2, 5, 3), dst=face_nodes)
-        return face_nodes
-
-    def get_grid_face_edges(self, grid: int, face_edges: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=(0, 1, 2, 3, 4, 5, 6, 1, 6, 7, 2), dst=face_edges)
-        return face_edges
-
-    def get_grid_nodes_per_face(
-        self, grid: int, nodes_per_face: np.ndarray
-    ) -> np.ndarray:
-        numpy.copyto(src=(4, 4, 3), dst=nodes_per_face)
-        return nodes_per_face
-
-    def get_grid_x(self, grid: int, x: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[0.0, 1.0, 2.0, 1.0, 3.0, 4.0], dst=x)
-        return x
-
-    def get_grid_y(self, grid: int, y: np.ndarray) -> np.ndarray:
-        numpy.copyto(src=[3.0, 1.0, 2.0, 4.0, 0.0, 3.0], dst=y)
-        return y
-
-    def get_grid_z(self, grid: int, z: np.ndarray) -> np.ndarray:
-        raise NotImplementedError("Do not know what z is")
-
-
-class DTypeModel(GridModel):
-    def __init__(self):
-        super().__init__()
-        self.dtype = numpy.dtype("float32")
-        self.value = numpy.array((1.1, 2.2, 3.3), dtype=self.dtype)
-
-    def get_var_type(self, name):
-        return str(self.dtype)
-
-    def get_var_itemsize(self, name):
-        return self.dtype.itemsize
-
-    def get_var_nbytes(self, name):
-        return self.dtype.itemsize * self.value.size
-
-    def get_value(self, name, dest):
-        numpy.copyto(src=self.value, dst=dest)
-        return dest
-
-    def get_value_at_indices(self, name, dest, inds):
-        numpy.copyto(src=self.value[inds], dst=dest)
-        return dest
-
-    def set_value(self, name, src):
-        self.value[:] = src
-
-    def set_value_at_indices(self, name, inds, src):
-        self.value[inds] = src
-
-
-class Float32Model(DTypeModel):
-    pass
-
-
-class Int32Model(DTypeModel):
-    def __init__(self):
-        super().__init__()
-        self.dtype = numpy.dtype("int32")
-        self.value = numpy.array((12, 24, 36), dtype=self.dtype)
-
-
-class BooleanModel(DTypeModel):
-    def __init__(self):
-        super().__init__()
-        self.dtype = numpy.dtype("bool")
-        self.value = numpy.array((True, False, True), dtype=self.dtype)
-
-
-class HugeModel(DTypeModel):
-    """Model which has value which does not fit in single message body
-
-    Can be run from command line with
-
-    ..code-block:: bash
-
-        run-bmi-server --path $PWD/test --name fake_models.HugeModel --port 55555 --debug
-    """
-
-    def __init__(self):
-        super().__init__()
-        self.dtype = numpy.dtype("float64")
-        # Create value which is bigger than 4Mb
-        dimension = (3 * GRPC_MAX_MESSAGE_LENGTH) // self.dtype.itemsize + 1000
-        self.value = numpy.ones((dimension,), dtype=self.dtype)
-
-
-class WithItemSizeZeroAndVarTypeFloat32Model(Float32Model):
-    def get_var_itemsize(self, name):
-        return 0
-
-
-class WithItemSizeZeroAndUnknownVarType(WithItemSizeZeroAndVarTypeFloat32Model):
-    def get_var_type(self, name):
-        return "real"
diff --git a/tests/src/models/test_abstract.py b/tests/src/models/test_abstract.py
deleted file mode 100644
index a06d82f3..00000000
--- a/tests/src/models/test_abstract.py
+++ /dev/null
@@ -1,339 +0,0 @@
-import logging
-import weakref
-from datetime import datetime, timezone
-from pathlib import Path
-from typing import Any, Iterable, Optional, Tuple
-from unittest.mock import patch
-
-import numpy as np
-import pytest
-import xarray as xr
-from bmipy import Bmi
-from grpc4bmi.bmi_optionaldest import OptionalDestBmi
-from numpy.testing import assert_array_equal
-from pydantic.config import ConfigDict
-from typing_extensions import Unpack
-
-from ewatercycle import CFG
-from ewatercycle.base.model import eWaterCycleModel
-from ewatercycle.base.parameter_set import ParameterSet
-
-pytest.skip(allow_module_level=True)
-
-
-@pytest.fixture
-def setup_config(tmp_path: Path):
-    CFG.parameterset_dir = tmp_path
-    config_file = tmp_path / "ewatercycle.yaml"
-    config_file.write_text("")
-    CFG.ewatercycle_config = config_file
-    yield CFG
-    CFG.ewatercycle_config = None
-    CFG.reload()
-
-
-class MockedModel(eWaterCycleModel):
-    bmi: OptionalDestBmi
-
-    def _make_bmi_instance(self) -> OptionalDestBmi:
-        return self.bmi
-
-    @property
-    def version(self) -> str:
-        return "0.4.2"
-
-    def setup(self, *args, **kwargs) -> Tuple[str, str]:
-        if "bmi" in kwargs:
-            # sub-class of AbstractModel should construct bmi
-            # using grpc4bmi Docker or Apptainer client
-            self.bmi = kwargs["bmi"]
-        return "foobar.cfg", "."
-
-    def get_value_as_xarray(self, name: str) -> xr.DataArray:
-        return xr.DataArray(
-            data=[[1.0, 2.0], [3.0, 4.0]],
-            coords={
-                "latitude": [42.25, 42.21],
-                "longitude": [-99.83, -99.32],
-                "time": "2014-09-06",
-            },
-            dims=["longitude", "latitude"],
-            name="Temperature",
-            attrs=dict(units="degC"),
-        )
-
-    def _coords_to_indices(
-        self, name: str, lat: Iterable[float], lon: Iterable[float]
-    ) -> Iterable[int]:
-        return np.array([0])
-
-    @property
-    def parameters(self) -> dict[str, Any]:
-        return {"area": 42}
-
-
-@pytest.fixture
-@patch("bmipy.Bmi")
-def bmi(MockedBmi):
-    mocked_bmi = MockedBmi()
-    mocked_bmi.get_start_time.return_value = 42.0
-    mocked_bmi.get_current_time.return_value = 43.0
-    mocked_bmi.get_end_time.return_value = 44.0
-    mocked_bmi.get_time_units.return_value = "days since 1970-01-01 00:00:00.0 00:00"
-    mocked_bmi.get_var_type.return_value = "float64"
-    mocked_bmi.get_var_itemsize.return_value = np.float64().size
-    mocked_bmi.get_var_nbytes.return_value = np.float64().size * 4
-    return mocked_bmi
-
-
-@pytest.fixture
-def model(bmi: Bmi):
-    mocked_model = MockedModel(bmi=bmi)
-    mocked_model.setup()
-    return mocked_model
-
-
-def test_construct():
-    with pytest.raises(TypeError) as excinfo:
-        eWaterCycleModel()
-    msg = str(excinfo.value)
-    assert "Can't instantiate abstract class" in msg
-    assert "_make_bmi_instance" in msg
-
-
-def test_setup(model):
-    result = model.setup()
-
-    expected = "foobar.cfg", "."
-    assert result == expected
-
-
-def test_initialize(model: MockedModel, bmi):
-    config_file = "foobar.cfg"
-    model.initialize(config_file)
-
-    bmi.initialize.assert_called_once_with(config_file)
-
-
-def test_finalize(model: MockedModel, bmi):
-    model.finalize()
-
-    bmi.finalize.assert_called_once_with()
-
-
-def test_finalize_resets_bmi(model: MockedModel, bmi):
-    model.finalize()
-
-    with pytest.raises(AttributeError) as excinfo:
-        model.bmi
-
-    assert "has no attribute 'bmi'" in str(excinfo.value)
-
-
-def test_update(model: MockedModel, bmi):
-    model.update()
-
-    bmi.update.assert_called_once_with()
-
-
-def test_get_value(bmi, model: MockedModel):
-    expected = np.array([1.0, 2.0, 3.0, 4.0])
-    bmi.get_value.return_value = expected
-
-    value = model.get_value("discharge")
-
-    assert_array_equal(value, expected)
-
-
-def test_get_value_at_coords(bmi, model: MockedModel):
-    expected = np.array([1.0])
-    bmi.get_value_at_indices.return_value = expected
-
-    value = model.get_value_at_coords("discharge", [-99.83], [42.25])
-
-    assert_array_equal(value, expected)
-
-
-def test_set_value(model: MockedModel, bmi):
-    value = np.array([1.0, 2.0])
-    model.set_value("precipitation", value)
-
-    bmi.set_value.assert_called_once_with("precipitation", value)
-
-
-def test_set_value_at_coords(model: MockedModel, bmi):
-    value = np.array([1.0])
-    model.set_value_at_coords("precipitation", [-99.83], [42.25], value)
-
-    bmi.set_value_at_indices.assert_called_once_with("precipitation", [0], value)
-
-
-def test_start_time(model: MockedModel):
-    time = model.start_time
-
-    assert time == pytest.approx(42.0)
-
-
-def test_end_time(model: MockedModel):
-    time = model.end_time
-
-    assert time == pytest.approx(44.0)
-
-
-def test_time(model: MockedModel):
-    time = model.time
-
-    assert time == pytest.approx(43.0)
-
-
-def test_time_units(model: MockedModel):
-    units = model.time_units
-
-    assert units == "days since 1970-01-01 00:00:00.0 00:00"
-
-
-def test_time_step(bmi, model: MockedModel):
-    bmi.get_time_step.return_value = 1.0
-
-    step = model.time_step
-
-    assert step == pytest.approx(1.0)
-
-
-def test_output_var_names(bmi, model: MockedModel):
-    bmi.get_output_var_names.return_value = ("discharge",)
-
-    names = model.output_var_names
-
-    assert names == ("discharge",)
-
-
-def test_get_value_as_xarray(model: MockedModel):
-    expected = xr.DataArray(
-        data=[[1.0, 2.0], [3.0, 4.0]],
-        coords={
-            "latitude": [42.25, 42.21],
-            "longitude": [-99.83, -99.32],
-            "time": "2014-09-06",
-        },
-        dims=["longitude", "latitude"],
-        name="Temperature",
-        attrs=dict(units="degC"),
-    )
-
-    dataarray = model.get_value_as_xarray("Temperature")
-
-    xr.testing.assert_equal(dataarray, expected)
-
-
-def start_time_as_isostr(model: MockedModel):
-    actual = model.start_time_as_isostr
-
-    expected = "1970-02-12T00:00:00Z"
-    assert expected == actual
-
-
-def end_time_as_isostr(model: MockedModel):
-    actual = model.end_time_as_isostr
-
-    expected = "1970-02-14T00:00:00Z"
-    assert expected == actual
-
-
-def time_as_isostr(model: MockedModel):
-    actual = model.time_as_isostr
-
-    expected = "1970-02-13T00:00:00Z"
-    assert expected == actual
-
-
-def start_time_as_datetime(model: MockedModel):
-    actual = model.start_time_as_isostr
-
-    expected = datetime(1970, 2, 12, tzinfo=timezone.utc)
-    assert expected == actual
-
-
-def end_time_as_datetime(model: MockedModel):
-    actual = model.end_time_as_isostr
-
-    expected = datetime(1970, 2, 14, tzinfo=timezone.utc)
-    assert expected == actual
-
-
-def time_as_datetime(model: MockedModel):
-    actual = model.time_as_isostr
-
-    expected = datetime(1970, 2, 13, tzinfo=timezone.utc)
-    assert expected == actual
-
-
-def test_delete_model_resets_bmi():
-    class Object:
-        """Target for weakref finalizer."""
-
-        pass
-
-    # Cannot use the `model` fixture, it doesn't play well with weakref
-    thismodel = MockedModel()
-    thismodel.setup(bmi=Object())
-    bmiref = weakref.finalize(thismodel.bmi, print, "Bmi got destroyed")
-
-    assert bmiref.alive
-    del thismodel
-    assert not bmiref.alive
-
-
-class TestCheckParameterSet:
-    def test_correct_version(self, setup_config):
-        ps = ParameterSet(
-            name="justatest",
-            directory="justatest",
-            config="justatest/config.ini",
-            target_model="mockedmodel",  # == lowered class name
-            supported_model_versions={"0.4.2"},
-        )
-        m = MockedModel(parameter_set=ps)
-        assert m.parameter_set == ps
-
-    def test_wrong_model(self, setup_config):
-        ps = ParameterSet(
-            name="justatest",
-            directory="justatest",
-            config="justatest/config.ini",
-            target_model="wrongmodel",
-            supported_model_versions={"0.4.2"},
-        )
-        with pytest.raises(ValueError) as excinfo:
-            MockedModel(parameter_set=ps)
-
-        expected = "Parameter set has wrong target model"
-        assert expected in str(excinfo.value)
-
-    def test_any_version(self, caplog, setup_config):
-        ps = ParameterSet(
-            name="justatest",
-            directory="justatest",
-            config="justatest/config.ini",
-            target_model="mockedmodel",  # == lowered class name
-            supported_model_versions=set(),
-        )
-        with caplog.at_level(logging.INFO):
-            MockedModel(parameter_set=ps)
-
-        expected = "is not explicitly listed in the supported model versions"
-        assert expected in caplog.text
-
-    def test_unsupported_version(self, setup_config):
-        ps = ParameterSet(
-            name="justatest",
-            directory="justatest",
-            config="justatest/config.ini",
-            target_model="mockedmodel",
-            supported_model_versions={"1.2.3"},
-        )
-        with pytest.raises(ValueError) as excinfo:
-            MockedModel(parameter_set=ps)
-
-        expected = "Parameter set is not compatible with version"
-        assert expected in str(excinfo.value)
diff --git a/tests/src/test_container.py b/tests/src/test_container.py
index 82bc7917..00cf5e15 100644
--- a/tests/src/test_container.py
+++ b/tests/src/test_container.py
@@ -14,7 +14,7 @@
     _parse_docker_url,
     start_container,
 )
-from tests.src.fake_models import Rect3DGridModel
+from ewatercycle.testing.fake_models import DummyModelWith2DRectilinearGrid
 
 images = [
     (
@@ -115,51 +115,45 @@ def eq(a, b):
     "orig_model, method_name, method_args, assert_func",
     [
         [
-            Rect3DGridModel(),
+            DummyModelWith2DRectilinearGrid(),
             "get_grid_shape",
-            (0, np.zeros((3,))),
-            lambda r: npeq(r, np.array([2, 3, 4])),
+            (0, np.zeros((2,))),
+            lambda r: npeq(r, np.array([3, 4])),
         ],
         [
-            Rect3DGridModel(),
+            DummyModelWith2DRectilinearGrid(),
             "get_grid_rank",
             (0,),
-            lambda r: eq(r, 3),
+            lambda r: eq(r, 2),
         ],
         [
-            Rect3DGridModel(),
+            DummyModelWith2DRectilinearGrid(),
             "get_grid_size",
             (0,),
-            lambda r: eq(r, 24),
+            lambda r: eq(r, 12),
         ],
         [
-            Rect3DGridModel(),
+            DummyModelWith2DRectilinearGrid(),
             "get_var_grid",
             ("air__temperature",),
             lambda r: eq(r, 0),
         ],
         [
-            Rect3DGridModel(),
+            DummyModelWith2DRectilinearGrid(),
             "get_grid_x",
             (0, np.empty((4,))),
             lambda r: npeq(r, np.array([0.1, 0.2, 0.3, 0.4])),
         ],
         [
-            Rect3DGridModel(),
+            DummyModelWith2DRectilinearGrid(),
             "get_grid_y",
             (0, np.empty((3,))),
             lambda r: npeq(r, np.array([1.1, 1.2, 1.3])),
         ],
-        [
-            Rect3DGridModel(),
-            "get_grid_z",
-            (0, np.empty((2,))),
-            lambda r: npeq(r, np.array([2.1, 2.2])),
-        ],
     ],
 )
 def test_bmi_proxy(orig_model, method_name, method_args, assert_func):
-    orig_model = Rect3DGridModel()
+    orig_model = DummyModelWith2DRectilinearGrid()
     model = BmiProxy(orig_model)
 
     method = getattr(model, method_name)
@@ -182,7 +176,7 @@ def force_apptainer(tmp_path: Path):
 @pytest.fixture
 def mock_bmi_client_apptainer():
     with mock.patch("ewatercycle.container.BmiClientApptainer") as mock_class:
-        mock_class.return_value = Rect3DGridModel()
+        mock_class.return_value = DummyModelWith2DRectilinearGrid()
         yield mock_class
 
 
@@ -194,7 +188,7 @@ def test_start_container(
 
     assert isinstance(container, OptionalDestBmi)
     assert isinstance(container.origin, MemoizedBmi)
-    assert isinstance(container.origin.origin, Rect3DGridModel)
+    assert isinstance(container.origin.origin, DummyModelWith2DRectilinearGrid)
     mock_bmi_client_apptainer.assert_called_once_with(
         image=force_apptainer.apptainer_filename,
         work_dir=str(tmp_path),
@@ -208,7 +202,7 @@ def test_start_container_without_wrapper(
     tmp_path: Path, force_apptainer, mock_bmi_client_apptainer: mock.MagicMock
 ):
     container = start_container(work_dir=tmp_path, image=force_apptainer, wrappers=[])
-    assert isinstance(container, Rect3DGridModel)
+    assert isinstance(container, DummyModelWith2DRectilinearGrid)
 
 
 def test_start_container_with_own_wrapper(
@@ -222,4 +216,4 @@ class MyWrapper(BmiProxy):
     )
 
     assert isinstance(container, MyWrapper)
-    assert isinstance(container.origin, Rect3DGridModel)
+    assert isinstance(container.origin, DummyModelWith2DRectilinearGrid)