11from asyncio import create_task , sleep
2- from unittest .mock import ANY , AsyncMock , MagicMock , patch
2+ from unittest .mock import ANY , AsyncMock , MagicMock , call , patch
33
44import pytest
5- from ophyd_async .testing import set_mock_value
5+ from ophyd_async .core import AsyncStatus
6+ from ophyd_async .testing import (
7+ callback_on_mock_put ,
8+ get_mock ,
9+ get_mock_put ,
10+ set_mock_value ,
11+ )
612
713from dodal .devices .robot import BartRobot , PinMounted , RobotLoadFailed , SampleLocation
814
@@ -28,8 +34,8 @@ async def _sleep(*args):
2834
2935 device ._load_pin_and_puck = AsyncMock (side_effect = _sleep )
3036
31- set_mock_value (device .error_code , (expected_error_code := 10 ))
32- set_mock_value (device .error_str , (expected_error_string := "BAD" ))
37+ set_mock_value (device .prog_error . code , (expected_error_code := 10 ))
38+ set_mock_value (device .prog_error . str , (expected_error_string := "BAD" ))
3339
3440 with pytest .raises (RobotLoadFailed ) as e :
3541 await device .set (SampleLocation (0 , 0 ))
@@ -82,32 +88,36 @@ async def test_given_program_not_running_and_pin_unmounting_but_new_pin_not_moun
8288 assert "Waiting on new pin loaded" in last_log
8389
8490
85- async def test_given_program_not_running_and_pin_unmounts_then_mounts_when_load_pin_then_pin_loaded () :
86- device = await _get_bart_robot ()
91+ async def set_with_happy_path ( device : BartRobot ) -> AsyncStatus :
92+ """Mocks the logic that the robot would do on a successful load"""
8793 device .LOAD_TIMEOUT = 0.05 # type: ignore
8894 set_mock_value (device .program_running , False )
8995 set_mock_value (device .gonio_pin_sensor , PinMounted .PIN_MOUNTED )
90- device .load = AsyncMock (side_effect = device .load )
9196 status = device .set (SampleLocation (15 , 10 ))
9297 await sleep (0.025 )
93- device .load .trigger .assert_called_once () # type:ignore
9498
9599 set_mock_value (device .gonio_pin_sensor , PinMounted .NO_PIN_MOUNTED )
96100 await sleep (0.025 )
97101
98102 set_mock_value (device .gonio_pin_sensor , PinMounted .PIN_MOUNTED )
103+ return status
104+
105+
106+ async def test_given_program_not_running_and_pin_unmounts_then_mounts_when_load_pin_then_pin_loaded ():
107+ device = await _get_bart_robot ()
108+ status = await set_with_happy_path (device )
99109 await status
100110 assert status .success
101111 assert (await device .next_puck .get_value ()) == 15
102112 assert (await device .next_pin .get_value ()) == 10
103- device .load . trigger . assert_called_once () # type:ignore
113+ get_mock_put ( device .load ). assert_called_once ()
104114
105115
106116async def test_given_waiting_for_pin_to_mount_when_no_pin_mounted_then_error_raised ():
107117 device = await _get_bart_robot ()
108118 status = create_task (device .pin_mounted_or_no_pin_found ())
109119 await sleep (0.2 )
110- set_mock_value (device .error_code , 25 )
120+ set_mock_value (device .prog_error . code , 25 )
111121 await sleep (0.01 )
112122 with pytest .raises (RobotLoadFailed ):
113123 await status
@@ -128,3 +138,42 @@ async def test_set_waits_for_both_timeouts(mock_wait_for: AsyncMock):
128138 device ._load_pin_and_puck = MagicMock ()
129139 await device .set (SampleLocation (1 , 2 ))
130140 mock_wait_for .assert_awaited_once_with (ANY , timeout = 0.02 )
141+
142+
143+ async def test_moving_the_robot_will_reset_error_if_light_curtain_is_tripped_and_still_throw_if_error_not_cleared ():
144+ device = await _get_bart_robot ()
145+ set_mock_value (device .controller_error .code , BartRobot .LIGHT_CURTAIN_TRIPPED )
146+
147+ with pytest .raises (RobotLoadFailed ) as e :
148+ await device .set (SampleLocation (1 , 2 ))
149+ assert e .value .error_code == 40
150+
151+ get_mock (device ).assert_has_calls (
152+ [
153+ call .reset .put (None , wait = True ),
154+ call .next_puck .put (ANY , wait = True ),
155+ call .next_pin .put (ANY , wait = True ),
156+ call .load .put (None , wait = True ),
157+ ]
158+ )
159+
160+
161+ async def test_moving_the_robot_will_reset_error_if_light_curtain_is_tripped_and_continue_if_error_cleared ():
162+ device = await _get_bart_robot ()
163+ set_mock_value (device .controller_error .code , BartRobot .LIGHT_CURTAIN_TRIPPED )
164+
165+ callback_on_mock_put (
166+ device .reset ,
167+ lambda * _ , ** __ : set_mock_value (device .controller_error .code , 0 ),
168+ )
169+
170+ await (await set_with_happy_path (device ))
171+
172+ get_mock (device ).assert_has_calls (
173+ [
174+ call .reset .put (None , wait = True ),
175+ call .next_puck .put (ANY , wait = True ),
176+ call .next_pin .put (ANY , wait = True ),
177+ call .load .put (None , wait = True ),
178+ ]
179+ )
0 commit comments