Issue #78 added tests for get frame with quaternions

tests/polaris_mocks.py

@@ -13,6 +13,15 @@
         "data/8700339.rom"]
     }
 
+SETTINGS_POLARIS_QUAT = {
+    "tracker type": "polaris",
+    "ports to probe": 20,
+    "romfiles" : [
+        "data/something_else.rom",
+        "data/8700339.rom"],
+    "use quaternions": "true"
+    }
+
 class MockPort:
     """A fake serial port for ndi"""
     device = 'bad port'
@@ -83,10 +92,10 @@ class MockBXFrameSource():
     def __init__(self):
         self.bx_frame_count = 0
         self.bx_call_count = 0
-        self.rotation = array([0, 0, 0, 0])
-        self.position = array([0, 0, 0])
-        self.velocity = array([10, -20, 5])
-        self.quality = array([1])
+        self.rotation = array([1., 0., 0., 0.])
+        self.position = array([0., 0., 0.])
+        self.velocity = array([10., -20., 5.])
+        self.quality = array([1.])
         self.tracked_tools = 0
 
     def setdevice(self, ndidevice):
@@ -111,18 +120,25 @@ def mockndiGetBXFrame(self, _device, port_handle): #pylint:disable=invalid-name
 
         return self.bx_frame_count
 
-    def mockndiGetBXTransform(self, _device, _port_handle): #pylint:disable=invalid-name
+    def mockndiGetBXTransform(self, _device, port_handle): #pylint:disable=invalid-name
         """
         Mock of ndiGetBXTransform. To enable a simple test of tracking
         smoothing translate the mock object between frames. Full testing of
         the averaging code is in the base class
         sksurgerycore.tests.algorithms
+        The base ndicapi library uses Py_BuildValue to return the transform
+        as a tuple of double float values, so we also
+        return a tuple
         """
         assert self.bx_frame_count > 0
-        #the base ndicapi library uses Py_BuildValue to return the transform
-        #as a tuple of double float values, so let's make sure we're als
-        #returning a tuple
-        return tuple(concatenate((self.rotation, self.position, self.quality)))
+        ph_int = int.from_bytes(port_handle, byteorder = 'little')
+        if ph_int == 0:
+            self.position = self.velocity * self.bx_frame_count
+            return tuple(concatenate((self.rotation, self.position,
+                self.quality)))
+
+        return tuple(concatenate((self.rotation, array([0, 0, 0]),
+            self.quality)))
 
     def mockndiGetBXTransformMissing(self, _device, _port_handle): #pylint:disable=invalid-name
         """Mock of ndiGetBXTransform"""

tests/test_sksurgerynditracker_mockndi_getframe.py

@@ -5,7 +5,8 @@
 import numpy as np
 from sksurgerynditracker.nditracker import NDITracker
 
-from tests.polaris_mocks import SETTINGS_POLARIS, mockndiProbe, \
+from tests.polaris_mocks import SETTINGS_POLARIS, SETTINGS_POLARIS_QUAT, \
+        mockndiProbe, \
         mockndiOpen, mockndiGetError, mockComports, \
         mockndiGetPHSRHandle, mockndiVER, \
         MockNDIDevice, MockBXFrameSource
@@ -66,10 +67,57 @@ def test_getframe_polaris_mock(mocker):
                                  [0.,0.,1.,10.],
                                  [0.,0.,0.,1.]])
     assert np.array_equal(expected_tracking_0, tracking[0])
-    expected_tracking_1 = np.array([[1.,0.,0.,0.],
-                                 [0.,1.,0.,0.],
-                                 [0.,0.,1.,0.],
-                                 [0.,0.,0.,1.]])
+    assert np.array_equal(expected_tracking_1, tracking[1])
+    assert tracking_quality.count(1.) == 2
+
+    del tracker
+
+def test_getframe_polaris_mock_quat(mocker):
+    """
+    Checks that get frame works with quaternions
+    """
+    tracker = None
+    bxsource = MockBXFrameSource()
+    ndidevice = MockNDIDevice()
+    mocker.patch('serial.tools.list_ports.comports', mockComports)
+    mocker.patch('ndicapy.ndiProbe', mockndiProbe)
+    mocker.patch('ndicapy.ndiOpen', mockndiOpen)
+    mocker.patch('ndicapy.ndiCommand', ndidevice.mockndiCommand)
+    mocker.patch('ndicapy.ndiGetError', mockndiGetError)
+    mocker.patch('ndicapy.ndiClose')
+    mocker.patch('ndicapy.ndiGetPHSRNumberOfHandles',
+            ndidevice.mockndiGetPHSRNumberOfHandles)
+    mocker.patch('ndicapy.ndiGetPHRQHandle', ndidevice.mockndiGetPHRQHandle)
+    mocker.patch('ndicapy.ndiPVWRFromFile')
+    mocker.patch('ndicapy.ndiGetPHSRHandle', mockndiGetPHSRHandle)
+    mocker.patch('ndicapy.ndiVER', mockndiVER)
+    mocker.patch('ndicapy.ndiGetBXFrame', bxsource.mockndiGetBXFrame)
+    mocker.patch('ndicapy.ndiGetBXTransform', bxsource.mockndiGetBXTransform)
+
+    tracker = NDITracker(SETTINGS_POLARIS_QUAT)
+
+    bxsource.setdevice(ndidevice)
+
+    (port_handles, time_stamps, frame_numbers, tracking,
+                tracking_quality ) = tracker.get_frame()
+
+    assert len(port_handles) == 2
+    assert len(time_stamps) == 2
+    assert frame_numbers.count(1) == 2
+    expected_tracking_0 = np.array([[1.,0.,0.,0.,10.,-20,5.]])
+    assert np.array_equal(expected_tracking_0, tracking[0])
+    expected_tracking_1 = np.array([[1.,0.,0.,0.,0.,0.,0.]])
+    assert np.array_equal(expected_tracking_1, tracking[1])
+    assert tracking_quality.count(1.) == 2
+
+    (port_handles, time_stamps, frame_numbers, tracking,
+                tracking_quality ) = tracker.get_frame()
+
+    assert len(port_handles) == 2
+    assert len(time_stamps) == 2
+    assert frame_numbers.count(2) == 2
+    expected_tracking_0 = np.array([[1.,0.,0.,0.,20.,-40,10.]])
+    assert np.array_equal(expected_tracking_0, tracking[0])
     assert np.array_equal(expected_tracking_1, tracking[1])
     assert tracking_quality.count(1.) == 2