Minor improvements

4 files changed, 91 insertions, 22 deletions

diff --git a/lab_control/frequency_response.py b/lab_control/frequency_response.py
index af399b1..8d46689 100644
--- a/lab_control/frequency_response.py
+++ b/lab_control/frequency_response.py
@@ -1,6 +1,6 @@
 from lab_control.function_generator import FunctionGenerator
 from lab_control.oscilloscope import Oscilloscope
-from lab_control.measurement import Measurement
+from lab_control.measurement import Measurement, getLinearRange
 
 class FrequencyResponseMeasurement(Measurement):
     def __init__(self):
@@ -13,13 +13,15 @@ class FrequencyResponseMeasurement(Measurement):
         self.data = None
 
     def measure(self, osc: Oscilloscope, fg: FunctionGenerator) -> None:
-        frequencyRange = self.maxFrequency - self.minFrequency
+        frequencies = getLinearRange(self.minFrequency, self.maxFrequency, self.steps)
         self.data = []
 
-        for i in range(0, self.steps):
-            frequency = self.minFrequency + i * frequencyRange / (self.steps - 1)
-            fg.setFrequency(self.functionGeneratorChannel, frequency)
+        for f in frequencies:
+            fg.setFrequency(self.functionGeneratorChannel, f)
             response = osc.measureAmplitude(self.oscilloscopeChannel)
-            self.data.append((frequency, response))
+            self.data.append((f, response))
 
         self.measurementDone = True
+
+    def saveToCSV(self, path: str) -> None:
+        pass
diff --git a/lab_control/measurement.py b/lab_control/measurement.py
index 78f393e..eca228d 100644
--- a/lab_control/measurement.py
+++ b/lab_control/measurement.py
@@ -1,6 +1,15 @@
 from lab_control.function_generator import FunctionGenerator
 from lab_control.oscilloscope import Oscilloscope
 
+def getLinearRange(first: float, last: float, steps: int) -> list[float]:
+    assert steps > 1, "Linear range requires at least two steps"
+    diff = last - first
+    stepSize = diff / (steps - 1)
+    return [first + i * stepSize for i in range(0, steps)]
+
 class Measurement:
     def measure(self, osc: Oscilloscope, fg: FunctionGenerator) -> None:
         pass
+
+    def saveToCSV(self, path: str) -> None:
+        pass
diff --git a/lab_control/test/frequency_response_measurement_test.py b/lab_control/test/frequency_response_test.py
index 5a73b3e..014b4f6 100644
--- a/lab_control/test/frequency_response_measurement_test.py
+++ b/lab_control/test/frequency_response_test.py
@@ -2,6 +2,7 @@ import pytest
 
 from lab_control.test.mock_lab import MockLab
 from lab_control.frequency_response import FrequencyResponseMeasurement
+from lab_control.measurement import getLinearRange
 
 @pytest.fixture
 def mockLab():
@@ -11,22 +12,9 @@ def mockLab():
 def uut(mockLab):
     return FrequencyResponseMeasurement()
 
-def test_frequencyResponseDefaults(uut):
-    assert uut.minFrequency == 20e0
-    assert uut.maxFrequency == 16e3
-    assert uut.steps == 50
-    assert uut.functionGeneratorChannel == 1
-    assert uut.oscilloscopeChannel == 1
-
-def test_frequencyResponseRamp(mockLab, uut):
-    uut.minFrequency = 100.0
-    uut.maxFrequency = 200.0
-    uut.steps = 11
-    uut.functionGeneratorChannel = 1
-    uut.oscilloscopeChannel = 1
-
+def prepareRampResponse(uut, mockLab):
     # Expect a ramp response from 0.5 to 1.5 * input amplitude
-    inputAmplitude = 1.0
+    inputAmplitude = 2.0
     minScale = 0.5
     maxScale = 1.5
 
@@ -40,7 +28,23 @@ def test_frequencyResponseRamp(mockLab, uut):
     mockLab.setAmplitude(uut.functionGeneratorChannel, inputAmplitude)
     mockLab.setOn(uut.functionGeneratorChannel)
 
-    expectedData = [(f, testFunction(f) * inputAmplitude) for f in [100.0, 110.0, 120.0, 130.0, 140.0, 150.0, 160.0, 170.0, 180.0, 190.0, 200.0]]
+    return [(f, testFunction(f) * inputAmplitude) for f in getLinearRange(uut.minFrequency, uut.maxFrequency, uut.steps)]
+
+def test_frequencyResponseDefaults(uut):
+    assert uut.minFrequency == 20e0
+    assert uut.maxFrequency == 16e3
+    assert uut.steps == 50
+    assert uut.functionGeneratorChannel == 1
+    assert uut.oscilloscopeChannel == 1
+
+def test_frequencyResponseRamp(mockLab, uut):
+    uut.minFrequency = 100.0
+    uut.maxFrequency = 200.0
+    uut.steps = 11
+    uut.functionGeneratorChannel = 1
+    uut.oscilloscopeChannel = 1
+
+    expectedData = prepareRampResponse(uut, mockLab)
 
     assert not uut.measurementDone
     assert uut.data == None
@@ -49,3 +53,4 @@ def test_frequencyResponseRamp(mockLab, uut):
 
     assert uut.measurementDone
     assert uut.data == expectedData
+
diff --git a/lab_control/test/measurement_test.py b/lab_control/test/measurement_test.py
new file mode 100644
index 0000000..07716a3
--- /dev/null
+++ b/lab_control/test/measurement_test.py
@@ -0,0 +1,53 @@
+import pytest
+
+from lab_control.measurement import *
+
+@pytest.fixture
+def uut():
+    return Measurement()
+
+def test_linearRangeAscending():
+    first = 100.0
+    last = 200.0
+    steps = 11
+
+    expectedRange = [
+        100.0, 110.0, 120.0, 130.0,
+        140.0, 150.0, 160.0, 170.0,
+        180.0, 190.0, 200.0
+    ]
+
+    actualRange = getLinearRange(first, last, steps)
+    assert actualRange == expectedRange
+
+def test_linearRangeDescending():
+    first = 200.0
+    last = 180.0
+    steps = 3
+
+    expectedRange = [
+        200.0, 190.0, 180.0
+    ]
+
+    actualRange = getLinearRange(first, last, steps)
+    assert actualRange == expectedRange
+
+def test_linearRangeTwoSteps():
+    first = 20.0
+    last = -20.0
+    steps = 2
+
+    expectedRange = [
+        20.0, -20.0
+    ]
+
+    actualRange = getLinearRange(first, last, steps)
+    assert actualRange == expectedRange
+
+def test_linearRangeOneStep():
+    first = 20.0
+    last = -20.0
+    steps = 1
+
+    with pytest.raises(AssertionError):
+        actualRange = getLinearRange(first, last, steps)