Update robot load unit test to check omega is still set

Author: rtuck99

tests/unit_tests/common/device_setup_plans/test_robot_load_unload.py

@@ -89,14 +89,27 @@ async def test_when_robot_unload_called_then_sample_area_prepared_before_load(
         and msg.args[0] == ApertureValue.OUT_OF_BEAM,
     )
 
-    assert_message_and_return_remaining(
+    msgs = assert_message_and_return_remaining(
+        msgs,
+        lambda msg: msg.command == "set"
+        and msg.obj is smargon.wrapped_omega.phase
+        and msg.args[0] == 0,
+    )
+
+    msgs = assert_message_and_return_remaining(
+        msgs,
+        lambda msg: msg.command == "wait"
+        and msg.kwargs["group"] == msgs[0].kwargs["group"],
+    )
+
+    msgs = assert_message_and_return_remaining(
         msgs,
         lambda msg: msg.command == "set"
         and msg.obj.name == "gonio"
         and msg.args[0] == CombinedMove(x=0, y=0, z=0, chi=0, phi=0),
     )
 
-    assert_message_and_return_remaining(
+    msgs = assert_message_and_return_remaining(
         msgs,
         lambda msg: msg.command == "set"
         and msg.obj is robot