@@ -127,6 +127,7 @@ def update_dispatch(sample: BaseDispatch, dispatch: BaseDispatch) -> BaseDispatc
127127 update_time = dispatch .update_time ,
128128 create_time = dispatch .create_time ,
129129 id = dispatch .id ,
130+ end_time = dispatch .end_time , # Ensure end_time is updated
130131 )
131132
132133
@@ -359,9 +360,12 @@ async def test_dispatch_schedule(
359360 type = "TEST_TYPE" ,
360361 )
361362 await test_env .client .create (** to_create_params (test_env .microgrid_id , sample ))
362- dispatch = Dispatch (test_env .client .dispatches (test_env .microgrid_id )[0 ])
363+ # Get the initial dispatch state from the client to use as a base for comparison
364+ initial_dispatch_from_client = Dispatch (
365+ test_env .client .dispatches (test_env .microgrid_id )[0 ]
366+ )
363367
364- next_run = dispatch .next_run_after (_now ())
368+ next_run = initial_dispatch_from_client .next_run_after (_now ())
365369 assert next_run is not None
366370
367371 fake_time .shift (next_run - _now () - timedelta (seconds = 1 ))
@@ -370,16 +374,21 @@ async def test_dispatch_schedule(
370374 # Expect notification of the dispatch being ready to run
371375 ready_dispatch = await test_env .running_state_change .receive ()
372376
373- assert ready_dispatch == dispatch
377+ # Use update_dispatch to create the expected object based on the initial state
378+ expected_ready_dispatch = Dispatch (update_dispatch (sample , ready_dispatch ))
379+ assert ready_dispatch == expected_ready_dispatch
374380
375- assert dispatch .duration is not None
381+ assert initial_dispatch_from_client .duration is not None
376382 # Shift time to the end of the dispatch
377- fake_time .shift (dispatch .duration + timedelta (seconds = 1 ))
383+ fake_time .shift (initial_dispatch_from_client .duration + timedelta (seconds = 1 ))
378384 await asyncio .sleep (1 )
379385
380386 # Expect notification to stop the dispatch
381387 done_dispatch = await test_env .running_state_change .receive ()
382- assert done_dispatch == dispatch
388+
389+ # Use update_dispatch again for the stop event comparison
390+ expected_done_dispatch = Dispatch (update_dispatch (sample , done_dispatch ))
391+ assert done_dispatch == expected_done_dispatch
383392
384393 await asyncio .sleep (1 )
385394
0 commit comments