improvements to creating mqtt clients; adding logic to handle if the BJ subclass's __init__ fails - we gracefully clean up

Author: CamDavidsonPilon

pioreactor/actions/pump.py

@@ -77,7 +77,7 @@ def __init__(
             self.calibration.hz,
             experiment=experiment,
             unit=unit,
-            pubsub_client=mqtt_client,
+            pub_client=mqtt_client,
             logger=logger,
         )
         self.pwm.lock()

pioreactor/automations/dosing/pid_morbidostat.py

@@ -55,6 +55,7 @@ def __init__(self, target_growth_rate: float | str, target_normalized_od: float
             experiment=self.experiment,
             job_name=self.job_name,
             target_name="growth_rate",
+            pub_client=self.pub_client,
         )
 
         assert isinstance(self.duration, float)

pioreactor/automations/temperature/thermostat.py

@@ -33,6 +33,7 @@ def __init__(self, target_temperature: float | str, **kwargs) -> None:
             job_name=self.job_name,
             target_name="temperature",
             output_limits=(-25, 25),  # avoid whiplashing
+            pub_client=self.pub_client,
         )
 
         self.set_target_temperature(target_temperature)
@@ -50,6 +51,10 @@ def _clamp_target_temperature(self, target_temperature: float) -> float:
 
         return clamp(0.0, target_temperature, self.MAX_TARGET_TEMP)
 
+    def on_disconnected(self) -> None:
+        super().on_disconnected()
+        self.pid.clean_up()
+
     def execute(self) -> UpdatedHeaterDC:
         while not hasattr(self, "pid"):
             # sometimes when initializing, this execute can run before the subclasses __init__ is resolved.

pioreactor/background_jobs/base.py

@@ -249,6 +249,22 @@ class _BackgroundJob(metaclass=PostInitCaller):
     # See pt.PublishableSetting type
     published_settings: dict[str, pt.PublishableSetting] = dict()
 
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        orig_init = cls.__init__
+
+        def wrapped_init(self, *args, **kwargs):
+            try:
+                orig_init(self, *args, **kwargs)
+            except Exception:
+                try:
+                    self.clean_up()
+                except Exception:
+                    pass
+                raise
+
+        cls.__init__ = wrapped_init
+
     def __init__(self, unit: str, experiment: str, source: str = "app") -> None:
         if self.job_name in DISALLOWED_JOB_NAMES:
             raise ValueError("Job name not allowed.")

pioreactor/background_jobs/stirring.py

@@ -242,7 +242,8 @@ def __init__(
             config.getfloat("stirring.config", "pwm_hz"),
             unit=self.unit,
             experiment=self.experiment,
-            pubsub_client=self.pub_client,
+            pub_client=self.pub_client,
+            logger=self.logger,
         )
         self.pwm.start(0)
         self.pwm.lock()
@@ -268,6 +269,7 @@ def __init__(
             job_name=self.job_name,
             target_name="rpm",
             output_limits=(-7.5, 7.5),  # avoid whiplashing
+            pub_client=self.pub_client,
         )
 
     def action_to_do_before_od_reading(self):
@@ -355,6 +357,8 @@ def on_disconnected(self) -> None:
             self.rpm_check_repeated_timer.cancel()
         with suppress(AttributeError):
             self.pwm.clean_up()
+        with suppress(AttributeError):
+            self.pid.clean_up()
         with suppress(AttributeError):
             if self.rpm_calculator:
                 self.rpm_calculator.clean_up()

pioreactor/background_jobs/temperature_automation.py

@@ -299,7 +299,14 @@ def setup_pwm(self) -> PWM:
         # impact (mainly: current sink), over the second. Ex: imagine freq=1hz, dc=40%, and the pump needs to run for
         # 0.3s. The influence of when the heat is one on the pump can be significant in a power-constrained system.
         pin = hardware.PWM_TO_PIN[hardware.HEATER_PWM_TO_PIN]
-        pwm = PWM(pin, hertz, unit=self.unit, experiment=self.experiment, pubsub_client=self.pub_client)
+        pwm = PWM(
+            pin,
+            hertz,
+            unit=self.unit,
+            experiment=self.experiment,
+            pub_client=self.pub_client,
+            logger=self.logger,
+        )
         pwm.start(0)
         return pwm
 

pioreactor/pubsub.py

@@ -346,7 +346,6 @@ def __init__(self, log_level: str, unit: str, experiment: str) -> None:
             allow_retained=False,
             client_id=f"{self.unit}_{self.experiment}_{self.log_level}_log_collector",
         )
-        assert self.client.is_connected
 
     def _collect_logs_into_bucket(self, message):
         # load the message

pioreactor/tests/test_automation_yamls.py

@@ -13,7 +13,9 @@
 
 
 def get_specific_yaml(path):
-    data = get(f"https://raw.githubusercontent.com/Pioreactor/pioreactorui/master/{path}")
+    r = get(f"https://raw.githubusercontent.com/Pioreactor/pioreactorui/master/{path}")
+    r.raise_for_status()
+    data = r.content
     return load(data.content, Loader=Loader)
 
 

pioreactor/tests/test_background_job.py

@@ -162,6 +162,32 @@ def update_state(msg: MQTTMessage) -> None:
     assert not is_pio_job_running("testjob")
 
 
+def test_what_happens_when_an_error_occurs_in_init_but_we_dont_catch() -> None:
+    class TestJob(BackgroundJob):
+        job_name = "testjob"
+
+        def __init__(self, unit: str, experiment: str) -> None:
+            super(TestJob, self).__init__(unit=unit, experiment=experiment)
+            raise ZeroDivisionError()
+
+    exp = "test_what_happens_when_an_error_occurs_in_init_but_we_dont_catch"
+    publish(f"pioreactor/unit/{exp}/testjob/$state", None, retain=True)
+    state = []
+
+    def update_state(msg: MQTTMessage) -> None:
+        state.append(msg.payload.decode())
+
+    subscribe_and_callback(update_state, f"pioreactor/unit/{exp}/testjob/$state")
+
+    with pytest.raises(ZeroDivisionError):
+        with TestJob(unit="unit", experiment=exp):
+            pass
+
+    pause()
+    assert state[-1] == "disconnected"
+    assert not is_pio_job_running("testjob")
+
+
 def test_state_transition_callbacks() -> None:
     class TestJob(BackgroundJob):
         job_name = "testjob"

pioreactor/tests/test_dosing_automation.py

@@ -1391,9 +1391,8 @@ def execute(self):
         assert len(bucket) >= 1
 
 
-@pytest.mark.xfail
 def test_a_failing_automation_cleans_duration_attr_in_mqtt_up() -> None:
-    experiment = "test_a_failing_automation_cleans_itself_up"
+    experiment = "test_a_failing_automation_cleans_duration_attr_in_mqtt_up"
 
     pubsub.publish(f"pioreactor/{get_unit_name()}/{experiment}/dosing_automation/duration", None, retain=True)