diff --git a/pioreactor/cli/calibrations.py b/pioreactor/cli/calibrations.py index c5ba8007..e8fc87a6 100644 --- a/pioreactor/cli/calibrations.py +++ b/pioreactor/cli/calibrations.py @@ -98,7 +98,7 @@ def run_calibration(ctx, device: str, protocol_name: str | None, y: bool) -> Non assistant = calibration_protocols.get((device, protocol_name)) if assistant is None: click.echo( - f"No protocols found for calibration device '{device}'. Available {device} protocols: {list(c[1] for c in calibration_protocols.keys() if c[0] == device)}" + f"No protocols found for device '{device}'. Available {device} protocols: {list(c[1] for c in calibration_protocols.keys() if c[0] == device)}" ) raise click.Abort() @@ -113,10 +113,10 @@ def run_calibration(ctx, device: str, protocol_name: str | None, y: bool) -> Non if not y: if click.confirm( - f"Do you want to set this calibration as the Active Calibration for {device}?", default=True + f"Do you want to set this calibration as the active calibration for {device}?", default=True ): calibration_struct.set_as_active_calibration_for_device(device) - click.echo(f"Set{calibration_struct.calibration_name} as the active calibration for {device}.") + click.echo(f"Set {calibration_struct.calibration_name} as the active calibration for {device}.") else: click.echo( f"Okay. You can use 'pio calibration set-active --device {device} --name {calibration_struct.calibration_name}' to set this calibration as the active one." @@ -127,6 +127,12 @@ def run_calibration(ctx, device: str, protocol_name: str | None, y: bool) -> Non ) +@calibration.command(name="protocols") +def list_protocols() -> None: + for device, protocol in calibration_protocols.keys(): + click.echo(f"{device}: {protocol}") + + @calibration.command(name="display") @click.option("--device", required=True, help="Calibration device.") @click.option("--name", "calibration_name", required=True, help="Name of calibration to display.") diff --git a/pioreactor/tests/test_calibrations.py b/pioreactor/tests/test_calibrations.py index 6488b32a..4b446352 100644 --- a/pioreactor/tests/test_calibrations.py +++ b/pioreactor/tests/test_calibrations.py @@ -10,6 +10,8 @@ from pioreactor import exc from pioreactor.calibrations import CALIBRATION_PATH +from pioreactor.calibrations import calibration_protocols +from pioreactor.calibrations import CalibrationProtocol from pioreactor.calibrations import load_active_calibration from pioreactor.calibrations import load_calibration from pioreactor.calibrations.utils import calculate_poly_curve_of_best_fit @@ -225,3 +227,27 @@ def test_mandys_data_for_pathological_poly() -> None: assert abs(mcal.predict(0.002) - curve_callable(0.002)) < 1e-10 assert abs(mcal.ipredict(0.002) - 0.002) < 0.1 + + +def test_custom_protocol(): + class CustomODCalibrationProtocol(CalibrationProtocol): + protocol_name = "custom" + target_device = "od" + + @staticmethod + def run(target_device, **kwargs): + pass + + assert calibration_protocols[("od", "custom")].__name__ == "CustomODCalibrationProtocol" + + class CustomCalibrationProtocolWithList(CalibrationProtocol): + protocol_name = "custom" + target_device = ["A", "B", "C"] + + @staticmethod + def run(target_device, **kwargs): + pass + + assert calibration_protocols[("A", "custom")].__name__ == "CustomCalibrationProtocolWithList" + assert calibration_protocols[("B", "custom")].__name__ == "CustomCalibrationProtocolWithList" + assert calibration_protocols[("C", "custom")].__name__ == "CustomCalibrationProtocolWithList"