Add module docstring and type hints

bessman · bessman · commit e18245d41e82 · 2020-11-05T21:05:34.000+01:00
diff --git a/tests/test_waveform_generator.py b/tests/test_waveform_generator.py
@@ -1,57 +1,63 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-Created on Fri Oct 23 13:55:56 2020
-
-@author: alexander
+"""Tests for PSL.waveform_generator.
+
+When integration testing, connect:
+    SQ1 -> LA1
+    SQ2 -> LA2
+    SQ3 -> LA3
+    SQ4 -> LA4
+    SI1 -> CH1
+    SI2 -> CH2
 """
 import time
 
 import numpy as np
-
 import pytest
+from scipy.optimize import curve_fit
+from _pytest.logging import LogCaptureFixture
+
 from PSL.logic_analyzer import LogicAnalyzer
 from PSL.oscilloscope import Oscilloscope
+from PSL.packet_handler import Handler
 from PSL.waveform_generator import PWMGenerator, WaveformGenerator
-from scipy.optimize import curve_fit
+
 
 MICROSECONDS = 1e-6
 RELTOL = 0.05
 GOOD_FIT = 0.99
 
 
-def r_squared(y, y_fit):
+def r_squared(y: np.ndarray, y_fit: np.ndarray) -> float:
     """Calculate the coefficient of determination."""
     ss_res = np.sum((y - y_fit) ** 2)  # Residual sum of squares.
     ss_tot = np.sum((y - y.mean()) ** 2)  # Total sum of squares.
     return 1 - (ss_res / ss_tot)
 
 
 @pytest.fixture
-def pwm(handler):
+def pwm(handler: Handler) -> PWMGenerator:
     handler._logging = True
     return PWMGenerator(handler)
 
 
 @pytest.fixture
-def wave(handler):
+def wave(handler: Handler) -> WaveformGenerator:
     handler._logging = True
     return WaveformGenerator(handler)
 
 
 @pytest.fixture
-def la(handler):
+def la(handler: Handler) -> LogicAnalyzer:
     handler._logging = True
     return LogicAnalyzer(handler)
 
 
 @pytest.fixture
-def scope(handler):
+def scope(handler: Handler) -> Oscilloscope:
     handler._logging = True
     return Oscilloscope(handler)
 
 
-def test_sine_wave(wave, scope):
+def test_sine_wave(wave: WaveformGenerator, scope: Oscilloscope):
     frequency = 500
     wave.load_equation("SI1", "sine")
     wave.generate("SI1", frequency)
@@ -78,7 +84,7 @@ def expected_f(x, amplitude, frequency, phase):
     assert coeff_of_det >= GOOD_FIT
 
 
-def test_triangle_wave(wave, scope):
+def test_triangle_wave(wave: WaveformGenerator, scope: Oscilloscope):
     frequency = 2000
     wave.load_equation("SI1", "tria")
     wave.generate("SI1", frequency)
@@ -107,7 +113,7 @@ def expected_f(x, amplitude, frequency, phase):
     assert coeff_of_det >= GOOD_FIT
 
 
-def test_superposition(wave, scope):
+def test_superposition(wave: WaveformGenerator, scope: Oscilloscope):
     frequency = 1000
     amplitude1 = 2
     amplitude2 = 1
@@ -147,7 +153,7 @@ def expected_f(x, amplitude1, amplitude2, frequency, phase):
     assert coeff_of_det >= GOOD_FIT
 
 
-def test_sine_phase(wave, scope):
+def test_sine_phase(wave: WaveformGenerator, scope: Oscilloscope):
     frequency = 500
     phase = 90
     wave.load_equation("SI1", "sine")
@@ -168,27 +174,27 @@ def expected_f(x, amplitude, frequency, phase):
     assert phase2_est - phase1_est == pytest.approx(phase * np.pi / 180, rel=RELTOL)
 
 
-def test_low_frequency_warning(caplog, wave):
+def test_low_frequency_warning(caplog: LogCaptureFixture, wave: WaveformGenerator):
     wave.generate("SI1", 1)
     assert "AC coupling" in caplog.text
 
 
-def test_low_frequency_error(wave):
+def test_low_frequency_error(wave: WaveformGenerator):
     with pytest.raises(ValueError):
         wave.generate("SI1", 0.05)
 
 
-def test_high_frequency_warning(caplog, wave):
+def test_high_frequency_warning(caplog: LogCaptureFixture, wave: WaveformGenerator):
     wave.generate("SI1", 1e4)
     assert "Frequencies above"
 
 
-def test_dimension_mismatch(wave):
+def test_dimension_mismatch(wave: WaveformGenerator):
     with pytest.raises(ValueError):
         wave.generate("SI2", [500, 1000])
 
 
-def test_pwm(pwm, la):
+def test_pwm(pwm: PWMGenerator, la: LogicAnalyzer):
     frequency = 5e4
     duty_cycle = 0.4
     pwm.generate("SQ1", frequency, duty_cycle)
@@ -198,7 +204,7 @@ def test_pwm(pwm, la):
     assert la.measure_duty_cycle("LA1")[1] == pytest.approx(duty_cycle, rel=RELTOL)
 
 
-def test_pwm_phase(pwm, la):
+def test_pwm_phase(pwm: PWMGenerator, la: LogicAnalyzer):
     frequency = 1e4
     duty_cycle = 0.5
     phase = 0.25
@@ -212,14 +218,14 @@ def test_pwm_phase(pwm, la):
     assert interval * MICROSECONDS == pytest.approx(frequency ** -1 * phase, rel=RELTOL)
 
 
-def test_set_state(pwm, la):
+def test_set_state(pwm: PWMGenerator, la: LogicAnalyzer):
     states = [True, False, True, True]
     pwm.set_state(*states)
     time.sleep(0.1)
     assert list(la.get_states().values()) == states
 
 
-def test_unchanged_state(pwm, la):
+def test_unchanged_state(pwm: PWMGenerator, la: LogicAnalyzer):
     frequency = 2.5e3
     duty_cycle = 0.9
     pwm.generate(["SQ1", "SQ4"], frequency, duty_cycle)
@@ -234,25 +240,25 @@ def test_unchanged_state(pwm, la):
     assert la.measure_duty_cycle("LA4")[1] == pytest.approx(duty_cycle, rel=RELTOL)
 
 
-def test_map_reference_clock(pwm, la):
+def test_map_reference_clock(pwm: PWMGenerator, la: LogicAnalyzer):
     prescaler = 5
     pwm.map_reference_clock(["SQ3"], prescaler)
     assert la.measure_frequency("LA3") == pytest.approx(
         128e6 / (1 << prescaler), rel=RELTOL
     )
 
 
-def test_pwm_high_frequency_error(pwm):
+def test_pwm_high_frequency_error(pwm: PWMGenerator):
     with pytest.raises(ValueError):
         pwm.generate("SQ1", 2e7, 0.5)
 
 
-def test_pwm_get_frequency(pwm):
+def test_pwm_get_frequency(pwm: PWMGenerator):
     frequency = 1500
     pwm.generate("SQ2", frequency, 0.1)
     assert pwm.frequency == pytest.approx(frequency, rel=RELTOL)
 
 
-def test_pwm_set_negative_frequency(pwm):
+def test_pwm_set_negative_frequency(pwm: PWMGenerator):
     with pytest.raises(ValueError):
         pwm.generate("SQ1", -1, 0.5)
