Fix tests

Author: codeSamuraii

tests/conftest.py

@@ -1,5 +1,4 @@
 import os
-import sys
 import time
 import httpx
 import pytest
@@ -13,6 +12,11 @@
 log = get_logger('setup-tests')
 
 
+@pytest.fixture(scope="session")
+def anyio_backend():
+    return 'asyncio'
+
+
 def find_free_port():
     """Find a free port on localhost."""
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -21,74 +25,100 @@ def find_free_port():
         port = s.getsockname()[1]
     return port
 
+def is_redis_reachable(tries: int = 5, delay: float = 1.0) -> bool:
+    """Check if Redis server is reachable."""
+    for _ in range(tries):
+        try:
+            redis_client.from_url('redis://127.0.0.1:6379').ping()
+            return True
+        except redis_client.ConnectionError:
+            time.sleep(delay)
+    return False
 
-@pytest.fixture(scope="session")
-def anyio_backend():
-    return 'asyncio'
-
+def is_uvicorn_reachable(port: int, tries: int = 5, delay: float = 1.0, *, base_url: str = 'http://127.0.0.1') -> bool:
+    """Check if Uvicorn server is reachable."""
+    for _ in range(tries):
+        try:
+            response = httpx.get(f'{base_url}:{port}/health', timeout=delay * 0.8)
+            if response.status_code == 200:
+                return True
+        except httpx.RequestError:
+            time.sleep(delay)
+    return False
 
-@pytest.fixture(scope="session")
-def live_server():
-    """Start uvicorn server in a subprocess."""
-    port = find_free_port()
-    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    processes = {}
+def start_redis_server(project_root: str) -> None:
+    print('\n', end='\r')
+    log.debug("- Starting Redis server...")
 
     try:
-        print()
-        log.debug(f"Starting test server...")
         redis_proc = subprocess.Popen(
             ['redis-server', '--port', '6379', '--save', '', '--appendonly', 'no'],
             cwd=project_root,
             stdout=subprocess.DEVNULL
         )
-        processes['redis'] = redis_proc
 
-        try:
-            time.sleep(2)
-            redis_client.from_url('redis://127.0.0.1:6379').ping()
-        except redis_client.ConnectionError:
-            log.error("Failed to connect to Redis server. Ensure Redis is running.")
-            raise
+        if not is_redis_reachable():
+            log.error("x Redis server did not start successfully.")
+            redis_proc.terminate()
+            raise RuntimeError("Could not start Redis server for tests.") from None
+
+        return redis_proc
 
+    except Exception as e:
+            log.error("x Failed to start Redis server.", exc_info=e)
+            redis_proc.kill()
+            raise RuntimeError("Could not start Redis server for tests.") from e
+
+def start_uvicorn_server(port: int, project_root: str) -> None:
+    try:
+        log.debug("- Starting uvicorn server...")
         uvicorn_proc = subprocess.Popen(
             ['uvicorn', 'app:app', '--host', '127.0.0.1', '--port', str(port)],
             cwd=project_root
         )
+
+        if not is_uvicorn_reachable(port):
+            log.error("x Uvicorn server did not start successfully.")
+            uvicorn_proc.terminate()
+            raise RuntimeError("Could not start Uvicorn server for tests.") from None
+
+        return uvicorn_proc
+
+    except Exception as e:
+        log.error("x Failed to start Uvicorn server.", exc_info=e)
+        uvicorn_proc.kill()
+        raise RuntimeError("Could not start Uvicorn server for tests.") from e
+
+@pytest.fixture(scope="session")
+def live_server():
+    """Start uvicorn server in a subprocess."""
+    port = find_free_port()
+    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    processes = {}
+
+    if not is_redis_reachable(tries=1):
+        redis_proc = start_redis_server(project_root)
+        processes['redis'] = redis_proc
+
+    if not is_uvicorn_reachable(port, tries=1):
+        uvicorn_proc = start_uvicorn_server(port, project_root)
         processes['uvicorn'] = uvicorn_proc
 
-        base_url = f'127.0.0.1:{port}'
-        max_retries = 5
-        for i in range(max_retries):
-            try:
-                response = httpx.get(f'http://{base_url}/health', timeout=5)
-                if response.status_code == 200:
-                    break
-            except Exception as e:
-                if i == max_retries - 1:
-                    uvicorn_proc.terminate()
-                    raise RuntimeError(f"Server failed to start after {max_retries} attempts") from None
-
-            time.sleep(2.0)
-
-        log.debug(f"Server started at {base_url}")
-        print()
-        yield base_url
+    yield f'127.0.0.1:{port}'
 
-        print()
-        for name, process in sorted(processes.items(), key=lambda x: -ord(x[0][0])):
-            if process.poll() is None:
-                log.debug(f"Terminating {name} process")
-                process.terminate()
-                try:
-                    process.wait(timeout=5)
-                except subprocess.TimeoutExpired: pass
-
-    finally:
-        for name, process in processes.items():
-            if process.poll() is None:
-                log.warning(f"Forcefully terminating {name}")
-                process.kill()
+    print()
+    for name in ['uvicorn', 'redis']:
+        process = processes.get(name)
+        if not process or process.poll() is not None:
+            continue
+
+        log.debug(f"- Terminating {name} process")
+        process.terminate()
+        try:
+            process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            log.warning(f"- {name} process did not terminate in time, killing it")
+            process.kill()
 
 
 @pytest.fixture

views/websockets.py

@@ -3,7 +3,6 @@
 import warnings
 from fastapi import WebSocket, APIRouter, WebSocketDisconnect, BackgroundTasks
 from fastapi.responses import PlainTextResponse
-from starlette.websockets import WebSocketClose
 from pydantic import ValidationError
 
 from lib.logging import get_logger
@@ -23,11 +22,9 @@ async def websocket_upload(websocket: WebSocket, uid: str):
     Then, the client must wait for the signal before sending file chunks.
     """
     if any(char not in string.ascii_letters + string.digits + '-' for char in uid):
-        await websocket.send_denial_response(PlainTextResponse(
-            "Invalid transfer ID. Must only contain alphanumeric characters and hyphens.",
-            status_code=400
-        ))
-        return WebSocketClose(code=1006, reason="Invalid transfer ID")
+        log.debug(f"△ Invalid transfer ID: {uid}")
+        await websocket.close(code=1008, reason="Invalid transfer ID")
+        return
 
     await websocket.accept()
     log.debug(f"△ Websocket upload request.")
@@ -64,10 +61,15 @@ async def websocket_upload(websocket: WebSocket, uid: str):
         log.warning("△ Receiver did not connect in time.")
         await websocket.send_text(f"Error: Receiver did not connect in time.")
         return
+    except Exception as e:
+        log.error("△ Error while waiting for receiver connection.", exc_info=e)
+        await websocket.send_text("Error: Error while waiting for receiver connection.")
+        return
 
-    transfer.info("△ Starting upload...")
+    transfer.debug("△ Sending go-ahead...")
     await websocket.send_text("Go for file chunks")
 
+    transfer.info("△ Starting upload...")
     await transfer.collect_upload(
         stream=websocket.iter_bytes(),
         on_error=send_error_and_close(websocket),