From 6f5895a1745e8f0a9b4e48312fba8af2a97e0fb4 Mon Sep 17 00:00:00 2001
From: Dan Allan <dallan@bnl.gov>
Date: Fri, 4 Oct 2024 17:13:41 -0400
Subject: [PATCH] Fix and test 'chunks' escape hatch.

---
 databroker/mongo_normalized.py  |  2 +-
 databroker/tests/test_broker.py | 47 +++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/databroker/mongo_normalized.py b/databroker/mongo_normalized.py
index 6e76f1433..90d6cb765 100644
--- a/databroker/mongo_normalized.py
+++ b/databroker/mongo_normalized.py
@@ -160,7 +160,7 @@ def structure_from_descriptor(descriptor, sub_dict, max_seq_num, unicode_columns
         numpy_dtype = dtype.to_numpy_dtype()
         if "chunks" in field_metadata:
             # If the Event Descriptor tells us a preferred chunking, use that.
-            suggested_chunks = field_metadata["chunks"]
+            suggested_chunks = tuple(tuple(chunks) for chunks in field_metadata["chunks"])
         elif (0 in shape) or (numpy_dtype.itemsize == 0):
             # special case to avoid warning from dask
             suggested_chunks = shape
diff --git a/databroker/tests/test_broker.py b/databroker/tests/test_broker.py
index e20bd3898..9a1d5f26d 100644
--- a/databroker/tests/test_broker.py
+++ b/databroker/tests/test_broker.py
@@ -1308,3 +1308,50 @@ def test_direct_img_read(db, RE, hw):
     c = db.v2
     uid, = get_uids(RE(count([hw.direct_img], 5)))
     c[uid]["primary"]["data"]["img"][:]
+
+
+def test_img_explicit_chunks(db, RE, hw, tmpdir):
+    "Test using explicit chunk size"
+    from ophyd import sim
+
+    RE.subscribe(db.insert)
+    if not hasattr(db, "v2"):
+        raise pytest.skip("v0 has no v2 accessor")
+    c = db.v2
+
+
+    class Detector1(sim.SynSignalWithRegistry):
+        def describe(self):
+            res = super().describe()
+            (key,) = res
+            shape = res[key]["shape"]
+            assert len(shape) == 2
+            res[key]["chunks"] = [[5], [1] * shape[0], [shape[1]]]
+            return res
+
+    class Detector2(sim.SynSignalWithRegistry):
+        def describe(self):
+            res = super().describe()
+            (key,) = res
+            shape = res[key]["shape"]
+            res[key]["chunks"] = [[5], [2] * (shape[0] // 2), [shape[1]]]
+            return res
+
+    img1 = Detector2(
+        func=lambda: np.array(np.ones((10, 10))),
+        name="img",
+        labels={"detectors"},
+        save_path=str(tmpdir),
+    )
+    img2 = Detector2(
+        func=lambda: np.array(np.ones((10, 10))),
+        name="img",
+        labels={"detectors"},
+        save_path=str(tmpdir),
+    )
+    uid, = get_uids(RE(count([img1], 5)))
+    c[uid]["primary"]["data"]["img"][:]
+    assert c[uid]["primary"]["data"]["img"].chunks[1] == tuple([2] * 5)
+    uid, = get_uids(RE(count([img2], 5)))
+    c[uid]["primary"]["data"]["img"][:]
+    assert c[uid]["primary"]["data"]["img"].chunks[1] == tuple([2] * 5)