rework nichols_test.py

bnavigator · bnavigator · commit 8ee73b94c44e · 2020-07-26T23:41:36.000+02:00
diff --git a/control/tests/nichols_test.py b/control/tests/nichols_test.py
@@ -1,34 +1,28 @@
-#!/usr/bin/env python
-#
-# nichols_test.py - test Nichols plot
-# RMM, 31 Mar 2011
+"""nichols_test.py - test Nichols plot
 
-import unittest
-import numpy as np
-from control.matlab import *
+RMM, 31 Mar 2011
+"""
 
-class TestStateSpace(unittest.TestCase):
-    """Tests for the Nichols plots."""
+import pytest
 
-    def setUp(self):
-        """Set up a system to test operations on."""
+from control import StateSpace, nichols_plot, nichols
 
-        A = [[-3., 4., 2.], [-1., -3., 0.], [2., 5., 3.]]
-        B = [[1.], [-3.], [-2.]]
-        C = [[4., 2., -3.]]
-        D = [[0.]]
 
-        self.sys = StateSpace(A, B, C, D)
+@pytest.fixture()
+def tsys():
+    """Set up a system to test operations on."""
+    A = [[-3., 4., 2.], [-1., -3., 0.], [2., 5., 3.]]
+    B = [[1.], [-3.], [-2.]]
+    C = [[4., 2., -3.]]
+    D = [[0.]]
+    return StateSpace(A, B, C, D)
 
-    def testNicholsPlain(self):
-        """Generate a Nichols plot."""
-        nichols(self.sys)
 
-    def testNgrid(self):
-        """Generate a Nichols plot."""
-        nichols(self.sys, grid=False)
-        ngrid()
+def test_nichols(tsys, mplcleanup):
+    """Generate a Nichols plot."""
+    nichols_plot(tsys)
 
 
-if __name__ == "__main__":
-    unittest.main()
+def test_nicholds_alias(tsys, mplcleanup):
+    """Test the control.nichols alias and the grid=False parameter"""
+    nichols(tsys, grid=False)
