Refine unit tests

examples/streaming_web/backend/app.py

@@ -4,6 +4,7 @@
 from contextlib import asynccontextmanager
 
 import socketio
+import uvicorn
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi_limiter import FastAPILimiter
@@ -65,12 +66,18 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None]:
 register_routes(app)
 register_sockets(sio, redis_manager)
 
-# Run the application using Uvicorn ASGI server
-if __name__ == '__main__':
-    import uvicorn
+
+def run_server():
+    """
+    Run the application using Uvicorn ASGI server.
+    """
     uvicorn.run(
-        'examples.streaming_web.app:sio_app',
+        'examples.streaming_web.backend.app:sio_app',
         host='127.0.0.1',
         port=8000,
         log_level='info',
     )
+
+
+if __name__ == '__main__':
+    run_server()

examples/streaming_web/backend/routes.py

@@ -161,13 +161,17 @@ async def websocket_stream(websocket: WebSocket, label: str, key: str) -> None:
                 # Send the latest frame and warnings to the client
                 await websocket.send_json(message)
             else:
-                # Send an error message if no new data is available
+                # If no new data is available, close the connection
                 await websocket.send_json({'error': 'No new data available'})
+                await websocket.close()
+                break
 
     except WebSocketDisconnect:
         print('WebSocket disconnected')
     except Exception as e:
         print(f"Unexpected error: {e}")
+        # Close the WebSocket connection on error
+        await websocket.close()
     finally:
         print('WebSocket connection closed')
 

src/model_fetcher.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
 
 import datetime
+import logging
 import time
 from pathlib import Path
 
 import requests
 import schedule
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
 class ModelFetcher:
     """
@@ -73,7 +78,7 @@ def download_and_save_model(
         self.local_dir.mkdir(parents=True, exist_ok=True)
         with open(local_file_path, 'wb') as f:
             f.write(model_file_content)
-        print(f"Model {model} successfully updated at {local_file_path}")
+        logger.info(f"Model {model} successfully updated at {local_file_path}")
 
     def request_new_model(self, model: str, last_update_time: str) -> None:
         """
@@ -99,14 +104,14 @@ def request_new_model(self, model: str, last_update_time: str) -> None:
                     model_file_content = bytes.fromhex(data['model_file'])
                     self.download_and_save_model(model, model_file_content)
                 else:
-                    print(f"Model {model} is already up to date.")
+                    logger.info(f"Model {model} is already up to date.")
             else:
-                print(
+                logger.error(
                     f"Failed to fetch model {model}. "
                     f"Server returned status code: {response.status_code}",
                 )
         except requests.exceptions.RequestException as e:
-            print(f"Error requesting model {model}: {e}")
+            logger.error(f"Error requesting model {model}: {e}")
 
     def update_all_models(self):
         """
@@ -117,11 +122,11 @@ def update_all_models(self):
         """
         for model in self.models:
             try:
-                print(f"Checking for updates for model {model}...")
+                logger.info(f"Checking for updates for model {model}...")
                 last_update_time = self.get_last_update_time(model)
                 self.request_new_model(model, last_update_time)
             except Exception as e:
-                print(f"Failed to update model {model}: {e}")
+                logger.error(f"Failed to update model {model}: {e}")
 
 
 # Schedule the task to run every hour
@@ -130,11 +135,14 @@ def schedule_task():
     updater.update_all_models()
 
 
+def run_scheduler_loop():
+    logger.info('Starting scheduled tasks. Press Ctrl+C to exit.')
+    while True:
+        schedule.run_pending()
+        time.sleep(1)
+
+
 if __name__ == '__main__':
     # Execute the scheduled task every hour
     schedule.every(1).hour.do(schedule_task)
-
-    print('Starting scheduled tasks. Press Ctrl+C to exit.')
-    while True:
-        schedule.run_pending()
-        time.sleep(1)  # Sleep for 1 second
+    run_scheduler_loop()