@@ -590,19 +590,28 @@ def expected_service_running(
590
590
assertion_output_folder : Path ,
591
591
) -> Generator [ServiceRunning , None , None ]:
592
592
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
+ )
599
599
)
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
+ )
600
612
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 )
606
615
elapsed_time = arrow .utcnow () - started
607
616
608
617
wait_for_service_endpoint_responding (
@@ -635,16 +644,27 @@ def wait_for_service_running(
635
644
In which case this will need further adjutment"""
636
645
637
646
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
+ )
644
653
)
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 )
648
668
elapsed_time = arrow .utcnow () - started
649
669
wait_for_service_endpoint_responding (
650
670
node_id ,
0 commit comments