@@ -29,63 +29,53 @@ class EchoWebsocketHandler(tornado.websocket.WebSocketHandler):
2929 def on_message (self , msg ):
3030 return self .write_message (msg )
3131
32- @contextmanager
33- def tornado_server (port : int = 8001 ):
34- ready_event2 = threading .Event ()
35- io_loop2 = None
36- app2 = None
37- io_thread2 = None
38-
39- def run_io_loop2 ():
40- nonlocal io_loop2 , app2 # Use nonlocal to modify outer scope variables
41- io_loop2 = tornado .ioloop .IOLoop ()
42- io_loop2 .make_current ()
43- app2 = tornado .web .Application ([(r"/" , EchoWebsocketHandler )])
44- app2 .listen (port )
45- ready_event2 .set ()
46- io_loop2 .start ()
47-
48- io_thread2 = threading .Thread (target = run_io_loop2 )
49- io_thread2 .start ()
50- ready_event2 .wait ()
51-
52- try :
53- yield
54- finally :
55- io_loop2 .add_callback (io_loop2 .stop )
56- if io_thread2 :
57- io_thread2 .join (timeout = 5 )
58- if io_thread2 .is_alive ():
59- raise RuntimeError ("IOLoop failed to stop" )
32+
33+ @contextmanager
34+ def create_tornado_server (port : int ):
35+ """Base context manager for creating a Tornado server in a thread"""
36+ ready_event = threading .Event ()
37+ io_loop = None
38+ app = None
39+ io_thread = None
40+
41+ def run_io_loop ():
42+ nonlocal io_loop , app
43+ io_loop = tornado .ioloop .IOLoop ()
44+ io_loop .make_current ()
45+ app = tornado .web .Application ([(r"/" , EchoWebsocketHandler )])
46+ app .listen (port )
47+ ready_event .set ()
48+ io_loop .start ()
49+
50+ io_thread = threading .Thread (target = run_io_loop )
51+ io_thread .start ()
52+ ready_event .wait ()
53+
54+ try :
55+ yield io_loop , app , io_thread
56+ finally :
57+ io_loop .add_callback (io_loop .stop )
58+ if io_thread :
59+ io_thread .join (timeout = 5 )
60+ if io_thread .is_alive ():
61+ raise RuntimeError ("IOLoop failed to stop" )
62+
63+
64+ @contextmanager
65+ def tornado_server (port : int = 8001 ):
66+ """Simplified context manager that uses the base implementation"""
67+ with create_tornado_server (port ) as (_io_loop , _app , _io_thread ):
68+ yield
6069
6170
62- @pytest .mark .skipif (not os .environ .get ("CSP_TEST_WEBSOCKET" ), reason = "Skipping websocket adapter tests" )
6371class TestWebsocket :
6472 @pytest .fixture (scope = "class" , autouse = True )
6573 def setup_tornado (self , request ):
66- # Create class-level attributes
67- request .cls .ready_event = threading .Event ()
68-
69- def run_io_loop ():
70- request .cls .io_loop = tornado .ioloop .IOLoop ()
71- request .cls .io_loop .make_current ()
72- request .cls .app = tornado .web .Application ([(r"/" , EchoWebsocketHandler )])
73- request .cls .app .listen (8000 )
74- request .cls .ready_event .set () # Signal that setup is complete
75- request .cls .io_loop .start ()
76-
77- request .cls .io_thread = threading .Thread (target = run_io_loop )
78- request .cls .io_thread .start ()
79- request .cls .ready_event .wait () # Wait for IOLoop to be ready
80-
81- # Teardown
82- yield
83-
84- request .cls .io_loop .add_callback (request .cls .io_loop .stop )
85- if request .cls .io_thread :
86- request .cls .io_thread .join (timeout = 5 ) # Add timeout to prevent hanging
87- if request .cls .io_thread .is_alive ():
88- raise RuntimeError ("IOLoop failed to stop" )
74+ with create_tornado_server (8000 ) as (io_loop , app , io_thread ):
75+ request .cls .io_loop = io_loop
76+ request .cls .app = app
77+ request .cls .io_thread = io_thread
78+ yield
8979
9080 def test_send_recv_msg (self ):
9181 @csp .node
@@ -122,6 +112,7 @@ def g():
122112 )
123113 if not send_payload_subscribe :
124114 # We send payload via the dummy send function
115+ # The 'on_connect_payload sends the result
125116 ws .send (csp .null_ts (object ), connection_request = csp .const (conn_request ))
126117 subscribe_connection_request = (
127118 ConnectionRequest (uri = "ws://localhost:8000/" , action = ActionType .CONNECT )
@@ -301,37 +292,40 @@ def g():
301292 assert msgs ["recv2" ][0 ][1 ].msg == "hey world from 8001"
302293 assert msgs ["recv2" ][0 ][1 ].uri == "ws://localhost:8001/"
303294
304- def test_unkown_host_graceful_shutdown (self ):
305- @csp .graph
306- def g ():
307- ws = WebsocketAdapterManager ("wss://localhost/" )
308- assert ws ._properties ["port" ] == "443"
309- csp .stop_engine (ws .status ())
310-
311- csp .run (g , starttime = datetime .now (pytz .UTC ), realtime = True )
312-
313- def test_send_recv_burst_json (self ):
295+ @pytest .mark .parametrize ("dynamic" , [False , True ])
296+ def test_send_recv_burst_json (self , dynamic ):
314297 class MsgStruct (csp .Struct ):
315298 a : int
316299 b : str
317300
318301 @csp .node
319- def send_msg_on_open (status : ts [Status ]) -> ts [str ]:
320- if csp .ticked (status ):
321- return MsgStruct (a = 1234 , b = "im a string" ).to_json ()
322-
323- @csp .node
324- def my_edge_that_handles_burst (objs : ts [List [MsgStruct ]]) -> ts [bool ]:
302+ def my_edge_that_handles_burst (objs : ts [List [MsgStruct ]]):
325303 if csp .ticked (objs ):
326- return True
304+ # Does nothing but makes sure it's not pruned
305+ ...
327306
328307 @csp .graph
329308 def g ():
330- ws = WebsocketAdapterManager ("ws://localhost:8000/" )
331- status = ws .status ()
332- ws .send (send_msg_on_open (status ))
333- recv = ws .subscribe (MsgStruct , JSONTextMessageMapper (), push_mode = csp .PushMode .BURST )
334- _ = my_edge_that_handles_burst (recv )
309+ if dynamic :
310+ ws = WebsocketAdapterManager (dynamic = True )
311+ wrapped_recv = ws .subscribe (
312+ MsgStruct ,
313+ JSONTextMessageMapper (),
314+ push_mode = csp .PushMode .BURST ,
315+ connection_request = csp .const (
316+ ConnectionRequest (
317+ uri = "ws://localhost:8000/" , on_connect_payload = MsgStruct (a = 1234 , b = "im a string" ).to_json ()
318+ )
319+ ),
320+ )
321+ recv = csp .apply (wrapped_recv , lambda vals : [v .msg for v in vals ], List [MsgStruct ])
322+ else :
323+ ws = WebsocketAdapterManager ("ws://localhost:8000/" )
324+ status = ws .status ()
325+ ws .send (csp .apply (status , lambda _x : MsgStruct (a = 1234 , b = "im a string" ).to_json (), str ))
326+ recv = ws .subscribe (MsgStruct , JSONTextMessageMapper (), push_mode = csp .PushMode .BURST )
327+
328+ my_edge_that_handles_burst (recv )
335329 csp .add_graph_output ("recv" , recv )
336330 csp .stop_engine (recv )
337331
@@ -341,3 +335,12 @@ def g():
341335 innerObj = obj [0 ]
342336 assert innerObj .a == 1234
343337 assert innerObj .b == "im a string"
338+
339+ def test_unkown_host_graceful_shutdown (self ):
340+ @csp .graph
341+ def g ():
342+ ws = WebsocketAdapterManager ("wss://localhost/" )
343+ assert ws ._properties ["port" ] == "443"
344+ csp .stop_engine (ws .status ())
345+
346+ csp .run (g , starttime = datetime .now (pytz .UTC ), realtime = True )
0 commit comments