From 373fe3dbe7e3363e292ab2e914e0e0461e04ffb5 Mon Sep 17 00:00:00 2001
From: Simon Renblad <srenblad@m-labs.hk>
Date: Fri, 30 Jun 2023 13:47:11 +0800
Subject: [PATCH] test_datasets: add metadata tests

---
 artiq/test/test_datasets.py | 33 +++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/artiq/test/test_datasets.py b/artiq/test/test_datasets.py
index 871568a2a..74e658323 100644
--- a/artiq/test/test_datasets.py
+++ b/artiq/test/test_datasets.py
@@ -16,6 +16,9 @@ class MockDatasetDB:
     def get(self, key):
         return self.data[key][1]
 
+    def get_metadata(self, key):
+        return self.data[key][2]
+
     def update(self, mod):
         # Copy mod before applying to avoid sharing references to objects
         # between this and the DatasetManager, which would lead to mods being
@@ -30,6 +33,9 @@ class TestExperiment(EnvExperiment):
     def get(self, key):
         return self.get_dataset(key)
 
+    def get_metadata(self, key):
+        return self.get_dataset_metadata(key)
+
     def set(self, key, value, **kwargs):
         self.set_dataset(key, value, **kwargs)
 
@@ -82,9 +88,9 @@ class ExperimentDatasetCase(unittest.TestCase):
     def test_append_broadcast(self):
         self.exp.set(KEY, [], broadcast=True)
         self.exp.append(KEY, 0)
-        self.assertEqual(self.dataset_db.data[KEY][1], [0])
+        self.assertEqual(self.dataset_db.get(KEY), [0])
         self.exp.append(KEY, 1)
-        self.assertEqual(self.dataset_db.data[KEY][1], [0, 1])
+        self.assertEqual(self.dataset_db.get(KEY), [0, 1])
 
     def test_append_array(self):
         for broadcast in (True, False):
@@ -103,3 +109,26 @@ class ExperimentDatasetCase(unittest.TestCase):
         with self.assertRaises(KeyError):
             self.exp.append(KEY, 0)
 
+    def test_set_dataset_metadata(self):
+        self.exp.set(KEY, 0, unit="kV", precision=2)
+        md = {"unit": "kV", "precision": 2}        
+        self.assertEqual(self.exp.get_metadata(KEY), md)
+    
+    def test_metadata_default(self):
+        self.exp.set(KEY, 0)
+        self.assertEqual(self.exp.get_metadata(KEY), {})
+
+    def test_metadata_scale(self):
+        self.exp.set(KEY, 0, scale=1000)
+        self.assertEqual(self.exp.get_metadata(KEY), {"scale": 1000})
+
+    def test_metadata_broadcast(self):
+        self.exp.set(KEY, 0, unit="kV", precision=2, broadcast=True)
+        md = {"unit": "kV", "precision": 2}        
+        self.assertEqual(self.dataset_db.get_metadata(KEY), md)
+
+    def test_metadata_broadcast_default(self):
+        self.exp.set(KEY, 0, broadcast=True)
+        self.assertEqual(self.dataset_db.get_metadata(KEY), {})
+
+