feat(websockets): implement WebSockets V3 with granular Redis Pub/Sub
- Introduced WS_Message_V3 standardized Pydantic model and WS_Manager_V3. - Implemented /v3/ws/ endpoint with granular Redis routing to solve "noisy neighbor" scaling issues. - Added presence tracking using Redis Sets for group coordination. - Comprehensive test suite added (unit and integration) covering models, manager, and routing logic. - Documentation: Created V3 Frontend WebSocket Guide and Project design spec. - Updated main Frontend API guide and tests README with new standards.
This commit is contained in:
Scott Idem
@@ -29,6 +29,8 @@ This directory contains the automated and manual test scripts for the Aether Fas
|
||||
| `test_unit_payload_sanitization.py` | **Primary Logic Test**: Verifies payload stripping and ID resolution. |
|
||||
| `test_unit_router_stripping.py` | Simulates automatic removal of random IDs during updates. |
|
||||
| `test_unit_schema_logic.py` | Verifies V3 schema metadata extraction logic with mocked DB rows. |
|
||||
| `test_unit_websockets_v3.py` | Unit tests for the V3 WebSocket manager and message models. |
|
||||
| `test_unit_websockets_v3_router.py` | Verifies the V3 WebSocket endpoint logic and message routing. |
|
||||
|
||||
### Integration Tests (`tests/integration/`)
|
||||
| Script | Description |
|
||||
@@ -70,3 +72,17 @@ This directory contains the automated and manual test scripts for the Aether Fas
|
||||
|
||||
### Path Requirements
|
||||
Always run test scripts from the **project root** directory. Most scripts include `sys.path.append(os.getcwd())` to ensure local imports work correctly.
|
||||
|
||||
---
|
||||
|
||||
## 💡 Best Practices & Reminders
|
||||
|
||||
1. **Check Before Creating**: Always check the **Script Inventory** above to see if a test for your logic already exists, or find a similar one to use as a reference/template.
|
||||
2. **Docker & Service Restarts**: Remember that the Aether Platform runs in Docker. If you modify core application code (e.g., in `app/`), you must restart the FastAPI service for changes to take effect:
|
||||
```bash
|
||||
docker restart aether_container_env-ae_api-2
|
||||
```
|
||||
(Note: Restarts are NOT necessary if you are only modifying the test scripts themselves).
|
||||
3. **Clean Up**: Clean up any temporary or debug files created during testing. However, **keep your test scripts**! Refactor them slightly for future use and clarity so they remain valuable assets for the project.
|
||||
4. **Stay Current**: Update this `README.md` when you add new tests or learn something that could help others. This is a living document; keep the **Script Inventory** and tips up to date.
|
||||
5. **Commit Often**: Don't forget to commit your working code and tests before moving on to the next task!
|
||||
|
||||
132
tests/integration/test_int_websockets_v3.py
Normal file
132
tests/integration/test_int_websockets_v3.py
Normal file
@@ -0,0 +1,132 @@
|
||||
import sys
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Add project root to path
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
# --- Robust Mocking BEFORE App Imports ---
|
||||
class MockSettings:
|
||||
def __init__(self):
|
||||
self.REDIS = {'server': 'localhost', 'port': 6379}
|
||||
self.DB = {
|
||||
'server': 'localhost',
|
||||
'port': 3306,
|
||||
'username': 'user',
|
||||
'password': 'pass',
|
||||
'database': 'db',
|
||||
'connect_timeout': 10,
|
||||
'pool_recycle': 3600
|
||||
}
|
||||
self.JWT_KEY = 'fake-key'
|
||||
self.AETHER_CFG = {'id': '0'}
|
||||
self.LOG_PATH = {'app': '/tmp/ae.log'}
|
||||
self.FILES_PATH = {'hosted_files_root': '/tmp', 'hosted_tmp_root': '/tmp'}
|
||||
self.ORIGINS_REGEX = '.*'
|
||||
self.ORIGINS = []
|
||||
|
||||
@property
|
||||
def SQLALCHEMY_DB_URI(self) -> str:
|
||||
return "mysql://user:pass@localhost:3306/db"
|
||||
|
||||
mock_settings = MockSettings()
|
||||
mock_config = MagicMock()
|
||||
mock_config.settings = mock_settings
|
||||
sys.modules["app.config"] = mock_config
|
||||
|
||||
# Mock DB related modules to prevent connection attempts at import time
|
||||
sys.modules["app.db_sql"] = MagicMock()
|
||||
sys.modules["app.lib_sql_core"] = MagicMock()
|
||||
sys.modules["app.db_connection"] = MagicMock()
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
from app.main import app
|
||||
|
||||
# Assume local Redis is running for integration testing
|
||||
client = TestClient(app)
|
||||
|
||||
def test_v3_websocket_communication():
|
||||
print("\n--- Testing V3 WebSocket: Group & Direct Communication ---")
|
||||
|
||||
group_id = "test_group_v3"
|
||||
client_a_id = "client_a"
|
||||
client_b_id = "client_b"
|
||||
|
||||
try:
|
||||
# 1. Connect both clients
|
||||
with client.websocket_connect(f"/v3/ws/group/{group_id}/client/{client_a_id}") as ws_a, \
|
||||
client.websocket_connect(f"/v3/ws/group/{group_id}/client/{client_b_id}") as ws_b:
|
||||
|
||||
print("Connected Client A and Client B.")
|
||||
|
||||
# --- Scenario A: Group Message ---
|
||||
print("\n[Scenario A] Client A sends a GROUP message...")
|
||||
msg_group = {
|
||||
"msg_type": "msg",
|
||||
"target": "group",
|
||||
"msg": "Hello Group!"
|
||||
}
|
||||
ws_a.send_json(msg_group)
|
||||
|
||||
resp_a = ws_a.receive_json()
|
||||
resp_b = ws_b.receive_json()
|
||||
|
||||
print(f"Client A received: {resp_a.get('msg')}")
|
||||
print(f"Client B received: {resp_b.get('msg')}")
|
||||
|
||||
assert resp_a["msg"] == "Hello Group!"
|
||||
assert resp_b["msg"] == "Hello Group!"
|
||||
assert resp_b["from_id"] == client_a_id
|
||||
print("✅ Group messaging verified.")
|
||||
|
||||
# --- Scenario B: Echo Message ---
|
||||
print("\n[Scenario B] Client A sends an ECHO message...")
|
||||
msg_echo = {
|
||||
"msg_type": "msg",
|
||||
"target": "echo",
|
||||
"msg": "Only for me"
|
||||
}
|
||||
ws_a.send_json(msg_echo)
|
||||
|
||||
resp_a_echo = ws_a.receive_json()
|
||||
print(f"Client A received: {resp_a_echo.get('msg')}")
|
||||
|
||||
assert resp_a_echo["msg"] == "Only for me"
|
||||
print("✅ Echo messaging verified.")
|
||||
|
||||
# --- Scenario C: Direct Message ---
|
||||
print("\n[Scenario C] Client A sends a DIRECT message to Client B...")
|
||||
msg_direct = {
|
||||
"msg_type": "cmd",
|
||||
"target": "direct",
|
||||
"to_id": client_b_id,
|
||||
"cmd": "RUN_TEST"
|
||||
}
|
||||
ws_a.send_json(msg_direct)
|
||||
|
||||
resp_b_direct = ws_b.receive_json()
|
||||
print(f"Client B received command: {resp_b_direct.get('cmd')}")
|
||||
|
||||
assert resp_b_direct["cmd"] == "RUN_TEST"
|
||||
assert resp_b_direct["from_id"] == client_a_id
|
||||
print("✅ Direct messaging verified.")
|
||||
|
||||
except ConnectionRefusedError:
|
||||
print("\n⚠️ Skipping test: Local Redis not found on port 6379.")
|
||||
except Exception as e:
|
||||
if "Connection refused" in str(e):
|
||||
print("\n⚠️ Skipping test: Local Redis not found on port 6379.")
|
||||
else:
|
||||
raise e
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
test_v3_websocket_communication()
|
||||
print("\n🎉 V3 WebSocket Integration Test Finished!")
|
||||
except Exception as e:
|
||||
print(f"\n❌ TEST FAILED: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
74
tests/unit/test_unit_websockets_v3.py
Normal file
74
tests/unit/test_unit_websockets_v3.py
Normal file
@@ -0,0 +1,74 @@
|
||||
import sys
|
||||
import os
|
||||
import asyncio
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, AsyncMock, patch
|
||||
from datetime import datetime, timezone
|
||||
|
||||
# Add project root to path
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
# Mock app.config BEFORE imports
|
||||
mock_config = MagicMock()
|
||||
mock_config.settings = MagicMock()
|
||||
mock_config.settings.REDIS = {'server': 'localhost', 'port': 6379}
|
||||
sys.modules["app.config"] = mock_config
|
||||
|
||||
from app.lib_websockets_v3 import WS_Message_V3, WS_Manager_V3
|
||||
|
||||
class TestWSV3Library(unittest.TestCase):
|
||||
|
||||
def test_message_model_validation(self):
|
||||
print("\n--- Testing WS_Message_V3 Validation ---")
|
||||
data = {
|
||||
"msg_type": "cmd",
|
||||
"target": "group",
|
||||
"from_id": "client_abc",
|
||||
"group_id": "group_123",
|
||||
"cmd": "RELOAD",
|
||||
"payload": {"force": True}
|
||||
}
|
||||
msg = WS_Message_V3(**data)
|
||||
self.assertEqual(msg.version, "3")
|
||||
self.assertEqual(msg.cmd, "RELOAD")
|
||||
self.assertTrue(isinstance(msg.sent_at.isoformat(), str))
|
||||
print("✅ Model validation passed.")
|
||||
|
||||
def test_channel_name_generation(self):
|
||||
print("\n--- Testing Channel Name Generation ---")
|
||||
manager = WS_Manager_V3()
|
||||
channels = manager.get_channel_names("client_abc", "group_123")
|
||||
|
||||
self.assertIn("ws:client:client_abc", channels)
|
||||
self.assertIn("ws:group:group_123", channels)
|
||||
self.assertIn("ws:broadcast", channels)
|
||||
print("✅ Channel name generation passed.")
|
||||
|
||||
@patch('redis.asyncio.Redis.from_url')
|
||||
def test_publish_routing(self, mock_redis_factory):
|
||||
print("\n--- Testing Publish Routing ---")
|
||||
mock_redis = AsyncMock()
|
||||
mock_redis_factory.return_value = mock_redis
|
||||
|
||||
manager = WS_Manager_V3()
|
||||
|
||||
async def run_test():
|
||||
# 1. Test Group Routing
|
||||
msg_group = WS_Message_V3(
|
||||
msg_type="msg", target="group", from_id="sender", group_id="target_group"
|
||||
)
|
||||
await manager.publish_message(msg_group)
|
||||
mock_redis.publish.assert_called_with("ws:group:target_group", unittest.mock.ANY)
|
||||
|
||||
# 2. Test Direct Routing
|
||||
msg_direct = WS_Message_V3(
|
||||
msg_type="msg", target="direct", from_id="sender", to_id="target_client"
|
||||
)
|
||||
await manager.publish_message(msg_direct)
|
||||
mock_redis.publish.assert_called_with("ws:client:target_client", unittest.mock.ANY)
|
||||
|
||||
asyncio.run(run_test())
|
||||
print("✅ Publish routing logic passed.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
126
tests/unit/test_unit_websockets_v3_router.py
Normal file
126
tests/unit/test_unit_websockets_v3_router.py
Normal file
@@ -0,0 +1,126 @@
|
||||
import sys
|
||||
import os
|
||||
import asyncio
|
||||
import json
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, AsyncMock, patch
|
||||
|
||||
# Add project root to path
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
# Mock app.config BEFORE imports to prevent attempt to load real settings
|
||||
mock_config = MagicMock()
|
||||
mock_config.settings = MagicMock()
|
||||
mock_config.settings.REDIS = {'server': 'localhost', 'port': 6379}
|
||||
sys.modules["app.config"] = mock_config
|
||||
|
||||
# Mock DB related modules to prevent circular imports or DB connection attempts
|
||||
sys.modules["app.db_sql"] = MagicMock()
|
||||
sys.modules["app.lib_sql_core"] = MagicMock()
|
||||
sys.modules["app.db_connection"] = MagicMock()
|
||||
|
||||
from app.routers.websockets_v3 import v3_ws_endpoint
|
||||
|
||||
class TestWSV3Router(unittest.TestCase):
|
||||
|
||||
@patch('app.routers.websockets_v3.ws_manager_v3')
|
||||
def test_v3_ws_endpoint_logic(self, mock_manager):
|
||||
"""
|
||||
Tests the core logic of the V3 WebSocket endpoint, ensuring
|
||||
Redis subscription and bidirectional message handling are initiated.
|
||||
"""
|
||||
# 1. Setup WebSocket Mock
|
||||
mock_ws = AsyncMock()
|
||||
|
||||
# 2. Setup Redis PubSub Mock
|
||||
mock_pubsub = MagicMock()
|
||||
mock_pubsub.subscribe = AsyncMock()
|
||||
mock_pubsub.unsubscribe = AsyncMock()
|
||||
mock_pubsub.close = AsyncMock()
|
||||
|
||||
mock_message = {
|
||||
'type': 'message',
|
||||
'data': json.dumps({
|
||||
"version": "3",
|
||||
"msg_type": "msg",
|
||||
"target": "group",
|
||||
"from_id": "other_client",
|
||||
"msg": "Hello from Redis",
|
||||
"payload": {},
|
||||
"sent_at": "2026-01-30T12:00:00Z"
|
||||
})
|
||||
}
|
||||
|
||||
# Signal to coordinate loops
|
||||
msg_delivered = asyncio.Event()
|
||||
|
||||
# Counters to break the 'while True' loops in the endpoint
|
||||
get_msg_count = 0
|
||||
recv_json_count = 0
|
||||
|
||||
async def mock_get_message(*args, **kwargs):
|
||||
nonlocal get_msg_count
|
||||
get_msg_count += 1
|
||||
if get_msg_count == 1:
|
||||
msg_delivered.set()
|
||||
return mock_message
|
||||
await asyncio.sleep(0.05)
|
||||
# Raise CancelledError to terminate the loop cleanly
|
||||
raise asyncio.CancelledError("Terminate sender loop")
|
||||
|
||||
mock_pubsub.get_message = mock_get_message
|
||||
|
||||
mock_redis = MagicMock()
|
||||
mock_redis.pubsub.return_value = mock_pubsub
|
||||
|
||||
# 3. Setup Manager Mock
|
||||
mock_manager.get_redis = AsyncMock(return_value=mock_redis)
|
||||
mock_manager.update_presence = AsyncMock()
|
||||
mock_manager.publish_message = AsyncMock()
|
||||
mock_manager.get_channel_names.return_value = ["ws:group:test"]
|
||||
|
||||
# Mock incoming websocket message
|
||||
async def mock_receive_json():
|
||||
nonlocal recv_json_count
|
||||
recv_json_count += 1
|
||||
if recv_json_count == 1:
|
||||
# Wait until the sender loop has processed the Redis message
|
||||
await msg_delivered.wait()
|
||||
return {
|
||||
"msg_type": "msg",
|
||||
"target": "group",
|
||||
"msg": "Client A saying hi"
|
||||
}
|
||||
await asyncio.sleep(0.05)
|
||||
# Raise CancelledError to terminate the loop cleanly
|
||||
raise asyncio.CancelledError("Terminate receiver loop")
|
||||
|
||||
mock_ws.receive_json.side_effect = mock_receive_json
|
||||
|
||||
# 4. Run the endpoint logic
|
||||
async def run_endpoint():
|
||||
try:
|
||||
# Execute endpoint with a short timeout
|
||||
await asyncio.wait_for(
|
||||
v3_ws_endpoint(mock_ws, "test_group", "client_a"),
|
||||
timeout=0.5
|
||||
)
|
||||
except (asyncio.TimeoutError, asyncio.CancelledError):
|
||||
pass
|
||||
except Exception as e:
|
||||
# Suppress our expected loop-termination messages
|
||||
if "Terminate" not in str(e):
|
||||
raise
|
||||
|
||||
asyncio.run(run_endpoint())
|
||||
|
||||
# 5. Verifications
|
||||
mock_ws.accept.assert_called_once()
|
||||
mock_manager.update_presence.assert_any_call("client_a", "test_group", online=True)
|
||||
|
||||
# Verify message from Redis was forwarded to WebSocket
|
||||
mock_ws.send_text.assert_called()
|
||||
print("✅ WebSocket Router unit logic verified.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user