1 files changed, 86 insertions, 0 deletions
diff --git a/lab_control/test/frequency_response_unittest.py b/lab_control/test/frequency_response_unittest.py
new file mode 100644
index 0000000..eb7d0d0
--- /dev/null
+++ b/lab_control/test/frequency_response_unittest.py
@@ -0,0 +1,86 @@
+import pytest
+
+from lab_control.test.mock.mock_lab import MockLab
+from lab_control.frequency_response import FrequencyResponseMeasurement
+from lab_control.measurement import getLinearRange
+from lab_control.function_generator import FunctionGenerator
+
+@pytest.fixture
+def mockLab():
+    return MockLab()
+
+@pytest.fixture
+def uut(mockLab):
+    return FrequencyResponseMeasurement()
+
+def prepareRampResponse(uut, mockLab):
+    # Expect a ramp response from 0.5 to 1.5 * input amplitude
+    minScale = 0.5
+    maxScale = 1.5
+
+    def testFunction(f: float) -> float:
+        assert f >= uut.minFrequency and f <= uut.maxFrequency
+        frequencyPu = (f - uut.minFrequency) / (uut.maxFrequency - uut.minFrequency)
+        return minScale + maxScale * frequencyPu
+
+    mockLab.connectChannels(uut.functionGeneratorChannel, uut.oscilloscopeChannel)
+    mockLab.setTestFunction(uut.oscilloscopeChannel, testFunction)
+
+    return [(f, testFunction(f) * uut.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.inputAmplitude == 0.1
+    assert uut.functionGeneratorChannel == 1
+    assert uut.oscilloscopeChannel == 1
+
+def test_frequencyResponseRamp(uut, mockLab):
+    uut.minFrequency = 100.0
+    uut.maxFrequency = 200.0
+    uut.steps = 11
+    uut.functionGeneratorChannel = 1
+    uut.oscilloscopeChannel = 1
+    uut.inputAmplitude = 2.0
+
+    fg = mockLab.getFunctionGeneratorChannel(uut.functionGeneratorChannel)
+    osc = mockLab.getOscilloscopeChannel(uut.oscilloscopeChannel)
+
+    expectedData = prepareRampResponse(uut, mockLab)
+
+    assert fg.amplitude is None
+    assert fg.function is None
+    assert not fg.on
+    assert osc.voltsPerDiv is None
+    assert not uut.measurementDone
+    assert uut.data is None
+    
+    uut.measure(mockLab, mockLab)
+
+    assert fg.amplitude == uut.inputAmplitude
+    assert fg.function == FunctionGenerator.SINE
+    assert fg.on
+    assert osc.voltsPerDiv == uut.inputAmplitude / (mockLab.getDivisionsDisplayed() - 2)
+    assert uut.measurementDone
+    assert uut.data == expectedData
+
+def test_dumpToCSV(uut, mockLab, tmp_path):
+    uut.minFrequency = 10.0
+    uut.maxFrequency = 30.0
+    uut.steps = 5
+    csv = tmp_path / "data.csv"
+
+    assert not csv.exists()
+    uut.dumpToCSV(csv)
+    assert not csv.exists()
+
+    expectedData = prepareRampResponse(uut, mockLab)
+    uut.measure(mockLab, mockLab)
+
+    uut.dumpToCSV(csv)
+    savedData = csv.read_text().split("\n")
+
+    assert savedData[0] == "Frequency (Hz),Amplitude (V)"
+    for i, l in enumerate(savedData[1:]):
+        assert l == f"{expectedData[i][0]},{expectedData[i][1]}"