refactor(api): only restore changed currents

api/src/opentrons/hardware_control/backends/ot3controller.py

@@ -932,14 +932,32 @@ async def set_hold_current(self, axis_currents: OT3AxisMap[float]) -> None:
             self._current_settings[axis].hold_current = current
 
     @asynccontextmanager
-    async def restore_current(self) -> AsyncIterator[None]:
-        """Save the current."""
-        old_current_settings = deepcopy(self._current_settings)
+    async def motor_current(
+        self,
+        run_currents: OT3AxisMap[float] = {},
+        hold_currents: OT3AxisMap[float] = {},
+    ) -> AsyncIterator[None]:
+        """Update and restore current."""
+        assert self._current_settings
+        old_settings = deepcopy(self._current_settings)
+        if run_currents:
+            await self.set_active_current(run_currents)
+        if hold_currents:
+            await self.set_hold_current(hold_currents)
         try:
             yield
         finally:
-            self._current_settings = old_current_settings
-            await self.set_default_currents()
+            if run_currents:
+                await self.set_active_current(
+                    {ax: old_settings[ax].run_current for ax in run_currents.keys()}
+                )
+            if hold_currents:
+                await self.set_hold_current(
+                    {ax: old_settings[ax].hold_current for ax in hold_currents.keys()}
+                )
+            if not run_currents and not hold_currents:
+                self._current_settings = old_settings
+                await self.set_default_currents()
 
     @asynccontextmanager
     async def restore_z_r_run_current(self) -> AsyncIterator[None]:

api/src/opentrons/hardware_control/backends/ot3simulator.py

@@ -540,7 +540,11 @@ async def set_active_current(self, axis_currents: OT3AxisMap[float]) -> None:
         return None
 
     @asynccontextmanager
-    async def restore_current(self) -> AsyncIterator[None]:
+    async def motor_current(
+        self,
+        run_currents: OT3AxisMap[float] = {},
+        hold_currents: OT3AxisMap[float] = {},
+    ) -> AsyncIterator[None]:
         """Save the current."""
         yield
 

api/src/opentrons/hardware_control/ot3api.py

@@ -1329,21 +1329,21 @@ async def _set_plunger_current_and_home(
                 moves = self._build_moves(
                     origin, target_pos, instr.config.plunger_homing_configurations.speed
                 )
-                async with self._backend.restore_current():
-                    await self._backend.set_active_current(
-                        {axis: instr.config.plunger_homing_configurations.current}
-                    )
+                async with self._backend.motor_current(
+                    run_currents={
+                        axis: instr.config.plunger_homing_configurations.current
+                    }
+                ):
                     await self._backend.move(
                         origin,
                         moves[0],
                         MoveStopCondition.none,
                     )
                     await self._backend.home([axis], self.gantry_load)
         else:
-            async with self._backend.restore_current():
-                await self._backend.set_active_current(
-                    {axis: instr.config.plunger_homing_configurations.current}
-                )
+            async with self._backend.motor_current(
+                run_currents={axis: instr.config.plunger_homing_configurations.current}
+            ):
                 await self._backend.home([axis], self.gantry_load)
 
     async def _retrieve_home_position(
@@ -1695,10 +1695,11 @@ async def _move_to_plunger_bottom(
         # NOTE: plunger position (mm) decreases up towards homing switch
         # NOTE: if already at BOTTOM, we still need to run backlash-compensation movement,
         #       because we do not know if we arrived at BOTTOM from above or below.
-        async with self._backend.restore_current():
-            await self._backend.set_active_current(
-                {pip_ax: instrument.config.plunger_homing_configurations.current}
-            )
+        async with self._backend.motor_current(
+            run_currents={
+                pip_ax: instrument.config.plunger_homing_configurations.current
+            }
+        ):
             if self._current_position[pip_ax] < backlash_pos[pip_ax]:
                 await self._move(
                     backlash_pos,
@@ -1880,10 +1881,9 @@ async def _force_pick_up_tip(
         self, mount: OT3Mount, pipette_spec: PickUpTipSpec
     ) -> None:
         for press in pipette_spec.presses:
-            async with self._backend.restore_current():
-                await self._backend.set_active_current(
-                    {axis: current for axis, current in press.current.items()}
-                )
+            async with self._backend.motor_current(
+                run_currents={axis: current for axis, current in press.current.items()}
+            ):
                 target_down = target_position_from_relative(
                     mount, press.relative_down, self._current_position
                 )
@@ -1899,10 +1899,11 @@ async def _force_pick_up_tip(
     async def _motor_pick_up_tip(
         self, mount: OT3Mount, pipette_spec: TipMotorPickUpTipSpec
     ) -> None:
-        async with self._backend.restore_current():
-            await self._backend.set_active_current(
-                {axis: current for axis, current in pipette_spec.currents.items()}
-            )
+        async with self._backend.motor_current(
+            run_currents={
+                axis: current for axis, current in pipette_spec.currents.items()
+            }
+        ):
             # Move to pick up position
             target_down = target_position_from_relative(
                 mount,

api/tests/opentrons/hardware_control/backends/test_ot3_controller.py

@@ -52,6 +52,7 @@
     TipStateType,
     FailedTipStateCheck,
     EstopState,
+    CurrentConfig,
 )
 from opentrons.hardware_control.errors import (
     FirmwareUpdateRequired,
@@ -1110,3 +1111,47 @@ async def test_requires_estop(
 
     with expectation:
         await controller.home([Axis.X, Axis.Y], gantry_load=GantryLoad.LOW_THROUGHPUT)
+
+
+@pytest.mark.parametrize(
+    "run_currents, hold_currents",
+    [
+        [{Axis.X: 1.0}, {}],
+        [{}, {Axis.X: 1.0}],
+        [{Axis.X: 1.0}, {Axis.X: 1.0}],
+        [{}, {}],
+    ],
+)
+async def test_motor_current(
+    controller: OT3Controller,
+    run_currents: OT3AxisMap[float],
+    hold_currents: OT3AxisMap[float],
+) -> None:
+    """Test that restore current actually works."""
+    controller._current_settings = {Axis.X: CurrentConfig(0.0, 0.0)}
+
+    with mock.patch.object(controller, "set_active_current") as mock_run_currents:
+        with mock.patch.object(controller, "set_hold_current") as mock_hold_currents:
+            with mock.patch.object(controller, "set_default_currents") as mock_default:
+
+                async with controller.motor_current(run_currents, hold_currents):
+                    await controller.update_position()
+
+                if not run_currents and not hold_currents:
+                    mock_default.assert_called_once()
+                    mock_run_currents.assert_not_called()
+                    mock_hold_currents.assert_not_called()
+                elif run_currents:
+                    mock_run_currents.assert_has_calls(
+                        [
+                            mock.call({Axis.X: 1.0}),
+                            mock.call({Axis.X: 0.0}),
+                        ],
+                    )
+                elif hold_currents:
+                    mock_hold_currents.assert_has_calls(
+                        [
+                            mock.call({Axis.X: 1.0}),
+                            mock.call({Axis.X: 0.0}),
+                        ],
+                    )

hardware-testing/hardware_testing/opentrons_api/helpers_ot3.py

@@ -581,10 +581,9 @@ async def move_plunger_absolute_ot3(
     if motor_current is None:
         await _move_coro
     else:
-        async with api._backend.restore_current():
-            await api._backend.set_active_current(
-                {Axis.of_main_tool_actuator(mount): motor_current}  # type: ignore[dict-item]
-            )
+        async with api._backend.motor_current(
+            run_currents={Axis.of_main_tool_actuator(mount): motor_current}
+        ):
             await _move_coro
 
 
@@ -613,10 +612,7 @@ async def move_tip_motor_relative_ot3(
     if motor_current is None:
         await _move_coro
     else:
-        async with api._backend.restore_current():
-            await api._backend.set_active_current(
-                {Axis.Q: motor_current}  # type: ignore[dict-item]
-            )
+        async with api._backend.motor_current(run_currents={Axis.Q: motor_current}):
             await _move_coro
 
 

hardware-testing/hardware_testing/production_qc/ninety_six_assembly_qc_ot3/test_droplets.py

@@ -152,8 +152,7 @@ async def _drop_tip(api: OT3API, trash: Point) -> None:
 async def _partial_pick_up_z_motion(
     api: OT3API, current: float, distance: float, speed: float
 ) -> None:
-    async with api._backend.restore_current():
-        await api._backend.set_active_current({Axis.Z_L: current})
+    async with api._backend.motor_current(run_currents={Axis.Z_L: current}):
         target_down = target_position_from_relative(
             OT3Mount.LEFT, Point(z=-distance), api._current_position
         )

hardware-testing/hardware_testing/production_qc/z_stage_qc_ot3.py

@@ -252,8 +252,9 @@ async def _force_gauge(
             force_output = []
             TH.start()
             try:
-                async with api._backend.restore_current():
-                    await api._backend.set_active_current({z_ax: test_current})
+                async with api._backend.motor_current(
+                    run_currents={z_ax: test_current}
+                ):
                     await api.move_to(
                         mount=mount,
                         abs_position=press_pos,