test: speed up tests

Author: AntoineRR

dev-requirements.txt

@@ -1,7 +1,7 @@
 -r requirements.txt
 flake8==4.0.1
 black==21.12b0
-websockets==10.1
+websocket-client==1.4.2
 maturin==0.12.11
 isort==5.10.1
 pre-commit==2.19.0

integration_tests/conftest.py

@@ -39,7 +39,7 @@ def session():
     base_routes = os.path.join(current_file_path, "./base_routes.py")
     command = ["python3", base_routes]
     process = spawn_process(command)
-    time.sleep(5)
+    time.sleep(1)
     yield
     kill_process(process)
 
@@ -50,7 +50,7 @@ def default_session():
     base_routes = os.path.join(current_file_path, "./base_routes.py")
     command = ["python3", base_routes]
     process = spawn_process(command)
-    time.sleep(5)
+    time.sleep(1)
     yield
     kill_process(process)
 
@@ -75,7 +75,7 @@ def dev_session():
     base_routes = os.path.join(current_file_path, "./base_routes.py")
     command = ["python3", base_routes, "--dev"]
     process = spawn_process(command)
-    time.sleep(5)
+    time.sleep(1)
     yield
     kill_process(process)
 
@@ -88,6 +88,6 @@ def test_session():
     base_routes = os.path.join(current_file_path, "./base_routes.py")
     command = ["python3", base_routes, "--dev"]
     process = spawn_process(command)
-    time.sleep(5)
+    time.sleep(1)
     yield
     kill_process(process)

integration_tests/test_web_sockets.py

@@ -1,19 +1,14 @@
-import asyncio
-
-from websockets import connect
+from websocket import create_connection
 
 BASE_URL = "ws://127.0.0.1:5000"
 
 
 def test_web_socket(session):
-    async def start_ws(uri):
-        async with connect(uri) as websocket:
-            assert await websocket.recv() == "Hello world, from ws"
-            await websocket.send("My name is?")
-            assert await websocket.recv() == "Whaaat??"
-            await websocket.send("My name is?")
-            assert await websocket.recv() == "Whooo??"
-            await websocket.send("My name is?")
-            assert await websocket.recv() == "*chika* *chika* Slim Shady."
-
-    asyncio.run(start_ws(f"{BASE_URL}/web_socket"))
+    ws = create_connection(f"{BASE_URL}/web_socket")
+    assert ws.recv() == "Hello world, from ws"
+    ws.send("My name is?")
+    assert ws.recv() == "Whaaat??"
+    ws.send("My name is?")
+    assert ws.recv() == "Whooo??"
+    ws.send("My name is?")
+    assert ws.recv() == "*chika* *chika* Slim Shady."

robyn/test-requirements.txt

@@ -4,5 +4,5 @@ watchdog
 requests==2.26.0
 uvloop==0.17.0
 multiprocess==0.70.12.2
-websockets==10.1
+websocket-client==1.4.2
 jinja2==3.0.2

src/server.rs

@@ -94,29 +94,27 @@ impl Server {
         let directories = self.directories.clone();
         let workers = Arc::new(workers);
 
-        let asyncio = py.import("asyncio").unwrap();
+        let asyncio = py.import("asyncio")?;
+        let event_loop = asyncio.call_method0("new_event_loop")?;
+        asyncio.call_method1("set_event_loop", (event_loop,))?;
 
-        let event_loop = asyncio.call_method0("new_event_loop").unwrap();
-        asyncio
-            .call_method1("set_event_loop", (event_loop,))
-            .unwrap();
         let startup_handler = self.startup_handler.clone();
         let shutdown_handler = self.shutdown_handler.clone();
 
-        let task_locals = Arc::new(pyo3_asyncio::TaskLocals::new(event_loop).copy_context(py)?);
-        let task_locals_cleanup = task_locals.clone();
+        let task_locals = pyo3_asyncio::TaskLocals::new(event_loop).copy_context(py)?;
+        let task_locals_copy = task_locals.clone();
 
         thread::spawn(move || {
             actix_web::rt::System::new().block_on(async move {
                 debug!("The number of workers are {}", workers.clone());
-                execute_event_handler(startup_handler, &task_locals)
+                execute_event_handler(startup_handler, &task_locals_copy)
                     .await
                     .unwrap();
 
                 HttpServer::new(move || {
                     let mut app = App::new();
 
-                    let task_locals = task_locals.clone();
+                    let task_locals = task_locals_copy.clone();
                     let directories = directories.read().unwrap();
 
                     // this loop matches three types of directory serving
@@ -168,7 +166,7 @@ impl Server {
                               global_request_headers,
                               body,
                               req| {
-                            pyo3_asyncio::tokio::scope_local((*task_locals).clone(), async move {
+                            pyo3_asyncio::tokio::scope_local(task_locals.clone(), async move {
                                 index(
                                     router,
                                     const_router,
@@ -198,13 +196,12 @@ impl Server {
             debug!("Ctrl c handler");
             Python::with_gil(|py| {
                 pyo3_asyncio::tokio::run(py, async move {
-                    execute_event_handler(shutdown_handler, &task_locals_cleanup)
+                    execute_event_handler(shutdown_handler, &task_locals.clone())
                         .await
                         .unwrap();
                     Ok(())
                 })
-                .unwrap();
-            })
+            })?
         }
         Ok(())
     }

src/web_socket_connection.rs

@@ -10,16 +10,13 @@ use pyo3_asyncio::TaskLocals;
 use uuid::Uuid;
 
 use std::collections::HashMap;
-use std::sync::Arc;
 
 /// Define HTTP actor
 #[derive(Clone)]
 struct MyWs {
     id: Uuid,
     router: HashMap<String, FunctionInfo>,
-    // can probably try removing arc from here
-    // and use clone_ref()
-    task_locals: Arc<TaskLocals>,
+    task_locals: TaskLocals,
 }
 
 fn get_function_output<'a>(
@@ -129,10 +126,8 @@ pub async fn start_web_socket(
     req: HttpRequest,
     stream: web::Payload,
     router: HashMap<String, FunctionInfo>,
-    task_locals: Arc<TaskLocals>,
+    task_locals: TaskLocals,
 ) -> Result<HttpResponse, Error> {
-    // execute the async function here
-
     ws::start(
         MyWs {
             router,