Implemented more tests and dsleep

ocs/ocs_twisted.py

@@ -1,6 +1,8 @@
 import threading
 from contextlib import contextmanager
 import time
+from autobahn.twisted.util import sleep as dsleep
+from twisted.internet.defer import inlineCallbacks
 
 
 class TimeoutLock:
@@ -156,7 +158,7 @@ class Pacemaker:
     Here is an example of how the Pacemaker can be used keep a 3 Hz quantized
     sample rate::
 
-        pm = Pacemaker(10, quantize=True)
+        pm = Pacemaker(3, quantize=True)
         take_data = True:
         while take_data:
             pm.sleep()
@@ -172,15 +174,29 @@ def __init__(self, sample_freq, quantize=False):
         if quantize and (sample_freq%1 != 0):
             raise ValueError("Quantization only works for frequencies that are whole numbers.")
 
-    def sleep(self):
-        """
-        Sleeps until the next calculated sampling time.
-        """
-        if time.time() < self.next_sample:
-            time.sleep(self.next_sample - time.time())
-
+    def _set_next_sample(self):
         self.next_sample = time.time() + self.sample_time
         if self.quantize:
             # Snaps "next_sample" to grid defined by sample_freq
             self.next_sample = (self.next_sample + self.sample_time/2) \
                                 // self.sample_time * self.sample_time
+
+    def sleep(self):
+        """
+        Sleeps until the next calculated sampling time.
+        """
+        now = time.time()
+        if now < self.next_sample:
+            time.sleep(self.next_sample - now)
+        self._set_next_sample()
+
+    @inlineCallbacks
+    def dsleep(self):
+        """
+        Sleeps in a non-blocking way by returning the deferred created by
+        twisted's sleep method.
+        """
+        now = time.time()
+        if now < self.next_sample:
+            yield dsleep(self.next_sample - now)
+        self._set_next_sample()

tests/test_pacemaker.py

@@ -1,8 +1,13 @@
-from ocs.ocs_twisted import Pacemaker
 import time
 import numpy as np
+import pytest
+from ocs.ocs_twisted import Pacemaker
 
-def test_pacemaker():
+def test_quantized():
+    """
+    Tests that Pacemaker forces regular sampling rate when quantize=True.
+
+    """
     sample_freq = 5
     pm = Pacemaker(sample_freq, quantize=True)
     times = []
@@ -13,4 +18,34 @@ def test_pacemaker():
         time.sleep(np.random.uniform(0, 1/sample_freq))
     diffs = np.diff(np.array(times))
     # Checks if the diffs (minus the first point due to quantization) match
-    assert np.all(np.abs(diffs - 1/sample_freq)[1:] < 1/sample_freq/5)
+    tolerance = 1/sample_freq / 5
+    assert np.all(np.abs(diffs - 1/sample_freq)[1:] < tolerance)
+
+
+def test_nonquantized():
+    """
+    Tests that pacemaker forces a regular sample rate when quantize=False.
+    In this case, all diffs should be less than the tolerance because
+    quantization doesn't mess up the first spacing.
+    """
+    sample_freq = 5
+    pm = Pacemaker(sample_freq, quantize=False)
+    times = []
+    for _ in range(10):
+        pm.sleep()
+        print("Sample time: {}".format(time.time()))
+        times.append(time.time())
+        time.sleep(np.random.uniform(0, 1/sample_freq))
+    diffs = np.diff(np.array(times))
+    # Checks if the diffs (minus the first point due to quantization) match
+    tolerance = 1/sample_freq / 5
+    assert np.all(np.abs(diffs - 1/sample_freq) < tolerance)
+
+
+def test_non_integer_quantization():
+    """
+    Trying to quantize with a non-integer sampling frequency should raise an
+    error.
+    """
+    with pytest.raises(ValueError):
+        Pacemaker(5.5, quantize=True)