Added set_data

pulse_lib/tests/mock_m4i.py

@@ -17,6 +17,8 @@ class MockM4i(Instrument):
         self.add_parameter('enable_channels', set_cmd=None, initial_value=0)
         self.add_parameter('data_memory_size', set_cmd=None, initial_value=0)
         self.add_parameter('pretrigger_memory_size', set_cmd=None, initial_value=16)
+        self._data = {}
+        self._n_triggers = 0
 
     def get_idn(self):
         return dict(vendor='Pulselib', model=type(self).__name__, serial='', firmware='')
@@ -41,18 +43,31 @@ class MockM4i(Instrument):
         pass
 
     def setup_multi_recording(self, seg_size, n_triggers, boxcar_average):
-        self.segment_size(math.ceil(seg_size/16)*16 + 16)
+        pretrigger = self.pretrigger_memory_size()
+        self.segment_size(math.ceil(seg_size/16)*16 + pretrigger)
         self.data_memory_size(self.segment_size()*n_triggers)
+        self._n_triggers = n_triggers
+
+    def set_data(self, ch_num, data):
+        self._data[ch_num] = data
 
     def get_data(self):
+        pretrigger = self.pretrigger_memory_size()
         channels = self.enable_channels()
-        n_ch = 0
-        while channels:
-            if channels & 1:
-                n_ch += 1
-            channels >>= 1
+        ch_nums = []
+        for ch_num in range(4):
+            if channels & (1 << ch_num):
+                ch_nums.append(ch_num)
+        n_ch = len(ch_nums)
         n_samples = self.data_memory_size()
-        res = np.empty((n_ch, n_samples))
+        res = np.full((n_ch, n_samples), np.nan)
         for i in range(n_ch):
-            res[i] = np.linspace(0, 1.0, n_samples)
+            data = self._data.get(ch_nums[i], None)
+            if data is None:
+                res[i] = np.linspace(0, 1.0, n_samples)
+            else:
+                mem_data = np.full((self._n_triggers, self.segment_size()), np.nan)
+                d = data.reshape(self._n_triggers, -1)
+                mem_data[:,pretrigger:pretrigger+d.shape[1]] = d
+                res[i] = mem_data.flatten()
         return res