From 742f02cb68e58b3d6f889cf50540b5863fa77cfe Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Fri, 27 Sep 2024 19:12:32 +0100
Subject: [PATCH] Simple test of StemAllometry

---
 pyrealm/demography/t_model_functions.py         | 13 +++++++++++++
 tests/unit/demography/test_t_model_functions.py | 15 +++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/pyrealm/demography/t_model_functions.py b/pyrealm/demography/t_model_functions.py
index 4c3a6234..65641dbc 100644
--- a/pyrealm/demography/t_model_functions.py
+++ b/pyrealm/demography/t_model_functions.py
@@ -6,6 +6,7 @@
 """  # noqa: D205
 
 from dataclasses import InitVar, dataclass, field
+from typing import ClassVar
 
 import numpy as np
 from numpy.typing import NDArray
@@ -625,6 +626,18 @@ class StemAllometry:
             allometry values.
     """
 
+    allometry_attrs: ClassVar[tuple[str, ...]] = (
+        "dbh",
+        "stem_height",
+        "crown_area",
+        "crown_fraction",
+        "stem_mass",
+        "foliage_mass",
+        "sapwood_mass",
+        "canopy_r0",
+        "canopy_z_max",
+    )
+
     # Init vars
     stem_traits: InitVar[Flora | StemTraits]
     at_dbh: InitVar[NDArray[np.float32]]
diff --git a/tests/unit/demography/test_t_model_functions.py b/tests/unit/demography/test_t_model_functions.py
index 19795831..3cc65e2e 100644
--- a/tests/unit/demography/test_t_model_functions.py
+++ b/tests/unit/demography/test_t_model_functions.py
@@ -704,3 +704,18 @@ def test_calculate_dbh_from_height_edge_cases():
 
     # Undefined entries
     assert np.all(np.isnan(dbh) == np.array([[0, 0], [0, 0], [0, 0], [1, 0], [1, 1]]))
+
+
+def test_StemAllometry(rtmodel_flora, rtmodel_data):
+    """Test the StemAllometry class."""
+
+    from pyrealm.demography.t_model_functions import StemAllometry
+
+    stem_allometry = StemAllometry(
+        stem_traits=rtmodel_flora, at_dbh=rtmodel_data["dbh"][:, [0]]
+    )
+
+    # Check the variables provided by the rtmodel implementation
+    to_check = set(stem_allometry.allometry_attrs) - set(["canopy_r0", "canopy_z_max"])
+    for var in to_check:
+        assert np.allclose(getattr(stem_allometry, var), rtmodel_data[var])