@@ -590,19 +590,28 @@ def expected_service_running(
590590 assertion_output_folder : Path ,
591591) -> Generator [ServiceRunning , None , None ]:
592592 started = arrow .utcnow ()
593- with log_context (
594- logging . INFO , msg = f"Waiting for node to run. Timeout: { timeout } "
595- ) as ctx :
596- waiter = SocketIONodeProgressCompleteWaiter (
597- node_id = node_id ,
598- logger = ctx . logger ,
593+ with contextlib . ExitStack () as stack :
594+ ctx = stack . enter_context (
595+ log_context (
596+ logging . INFO ,
597+ msg = f"Waiting for node to run. Timeout: { timeout } " ,
598+ )
599599 )
600+ if is_service_legacy :
601+ ctx .logger .info (
602+ "⚠️ Legacy service detected. We are skipping websocket messages in this case! ⚠️"
603+ )
604+ else :
605+ waiter = SocketIONodeProgressCompleteWaiter (
606+ node_id = node_id ,
607+ logger = ctx .logger ,
608+ )
609+ stack .enter_context (
610+ websocket .expect_event ("framereceived" , waiter , timeout = timeout )
611+ )
600612 service_running = ServiceRunning (iframe_locator = None )
601- with websocket .expect_event ("framereceived" , waiter , timeout = timeout ):
602- if press_start_button :
603- _trigger_service_start (page , node_id )
604-
605- yield service_running
613+ if press_start_button :
614+ _trigger_service_start (page , node_id )
606615 elapsed_time = arrow .utcnow () - started
607616
608617 wait_for_service_endpoint_responding (
@@ -635,16 +644,27 @@ def wait_for_service_running(
635644 In which case this will need further adjutment"""
636645
637646 started = arrow .utcnow ()
638- with log_context (
639- logging . INFO , msg = f"Waiting for node to run. Timeout: { timeout } "
640- ) as ctx :
641- waiter = SocketIONodeProgressCompleteWaiter (
642- node_id = node_id ,
643- logger = ctx . logger ,
647+ with contextlib . ExitStack () as stack :
648+ ctx = stack . enter_context (
649+ log_context (
650+ logging . INFO ,
651+ msg = f"Waiting for node to run. Timeout: { timeout } " ,
652+ )
644653 )
645- with websocket .expect_event ("framereceived" , waiter , timeout = timeout ):
646- if press_start_button :
647- _trigger_service_start (page , node_id )
654+ if is_service_legacy :
655+ ctx .logger .info (
656+ "⚠️ Legacy service detected. We are skipping websocket messages in this case! ⚠️"
657+ )
658+ else :
659+ waiter = SocketIONodeProgressCompleteWaiter (
660+ node_id = node_id ,
661+ logger = ctx .logger ,
662+ )
663+ stack .enter_context (
664+ websocket .expect_event ("framereceived" , waiter , timeout = timeout )
665+ )
666+ if press_start_button :
667+ _trigger_service_start (page , node_id )
648668 elapsed_time = arrow .utcnow () - started
649669 wait_for_service_endpoint_responding (
650670 node_id ,
0 commit comments