fix mk_level2_image for 3d shape

CHANGES.rst

@@ -1,3 +1,8 @@
+0.22.0 (2024-08-06)
+===================
+
+- Fix mk_level2_image utility for 3d shape. [#378]
+
 0.21.0 (2024-08-06)
 ===================
 

src/roman_datamodels/maker_utils/_datamodels.py

@@ -98,8 +98,8 @@ def mk_level2_image(*, shape=(4088, 4088), n_groups=8, filepath=None, **kwargs):
     roman_datamodels.stnode.WfiImage
     """
     if len(shape) > 2:
-        shape = shape[1:3]
         n_groups = shape[0]
+        shape = shape[1:3]
 
         warnings.warn(
             f"{MESSAGE} assuming the first entry is n_groups followed by y, x. The remaining is thrown out!", UserWarning

tests/test_maker_utils.py

@@ -204,3 +204,14 @@ def test_keyword_only(node_class):
             assert param.kind == inspect.Parameter.VAR_KEYWORD
         else:
             assert param.kind == inspect.Parameter.KEYWORD_ONLY
+
+
+@pytest.mark.filterwarnings("ignore:This function assumes shape is 2D")
+def test_mk_level2_image_shape():
+    """
+    Regression test for https://github.com/spacetelescope/roman_datamodels/issues/377
+    where n_groups was incorrect when provided a 3d shape
+    """
+    n = maker_utils.mk_level2_image(shape=(2, 3, 4))
+    assert n.amp33.shape == (2, 4096, 128)
+    assert n.data.shape == (3, 4)