From b74fdc0f7a29563b56618c109bebe972fe97f560 Mon Sep 17 00:00:00 2001
From: Tom Willemsen <tom.willemsen@diamond.ac.uk>
Date: Fri, 18 Aug 2023 11:45:16 +0100
Subject: [PATCH] Add test for NTNDArray-type PVA variables

---
 ophyd/v2/tests/test_epics.py   | 25 +++++++++++++++++++++++++
 ophyd/v2/tests/test_records.db | 21 ++++++++++++++++++++-
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/ophyd/v2/tests/test_epics.py b/ophyd/v2/tests/test_epics.py
index eda40d09c..aea1ca011 100644
--- a/ophyd/v2/tests/test_epics.py
+++ b/ophyd/v2/tests/test_epics.py
@@ -300,6 +300,31 @@ async def test_pva_table(ioc: IOC) -> None:
             q.close()
 
 
+async def test_pva_ntdarray(ioc: IOC):
+    if ioc.protocol == "ca":
+        # CA can't do ndarray
+        return
+    initial = np.zeros(4, np.int64)
+    put = np.ones_like(initial)
+
+    descriptor = dict(dtype="array", shape=[4])
+
+    for i, p in [(initial, put), (put, initial)]:
+        backend = await ioc.make_backend(npt.NDArray[np.int64], "ntndarray")
+        # Make a monitor queue that will monitor for updates
+        q = MonitorQueue(backend)
+        try:
+            # Check descriptor
+            assert dict(source=backend.source, **descriptor) == await backend.get_descriptor()
+            # Check initial value
+            await q.assert_updates(pytest.approx(i))
+            # Put to new value and check that
+            await backend.put(p)
+            await q.assert_updates(pytest.approx(p))
+        finally:
+            q.close()
+
+
 async def test_non_existant_errors(ioc: IOC):
     backend = await ioc.make_backend(str, "non-existant", connect=False)
     # Can't use asyncio.wait_for on python3.8 because of
diff --git a/ophyd/v2/tests/test_records.db b/ophyd/v2/tests/test_records.db
index 307ad5cd3..e0eacfdec 100644
--- a/ophyd/v2/tests/test_records.db
+++ b/ophyd/v2/tests/test_records.db
@@ -241,4 +241,23 @@ record(waveform, "$(P)table:enum")
             "": {"+type": "meta", "+channel": "VAL"}
         }
     })
-}
\ No newline at end of file
+}
+
+record(waveform, "$(P)ntndarray")
+{
+  field(FTVL, "INT64")
+  field(NELM, "4")
+  field(INP, {const:[0, 0, 0, 0]})
+  field(PINI, "YES")
+  info(Q:group, {
+        "$(P)ntndarray":{
+            +id:"epics:nt/NTNDArray:1.0",
+            "value":{
+                +type:"any",
+                +channel:"VAL",
+                +trigger:"*",
+            },
+            "": {+type:"meta", +channel:"VAL"}
+        }
+    })
+}