Set up locust (#12)

* clean up ws session handling, limit potential contacts list

* setup locust

README.md

@@ -36,3 +36,7 @@ TEST_MODE=yes PWDEBUG=1 poetry run pytest app.py --base-url http://localhost:500
 # Deploy
 
 Check [deploy](https://github.com/callmephilip/tinychat/tree/deploy) branch for an example deployment approach. Additional deps:
+
+# Run load tests
+
+`poetry run locust -f locustfile.py --headless`

app.py

@@ -380,7 +380,7 @@ def direct_channel_placeholders(self) -> List[ChannelPlaceholder]:
                         select channel from channel_member where member={self.member.id}
                     )
                 )
-            ) and id != {self.member.id}
+            ) and id != {self.member.id} LIMIT 10
         """)))
 
 @dataclass
@@ -581,7 +581,7 @@ def LandingLayout(*content) -> FT:
     )
 
 def Layout(*content, m: Member, w: Workspace, channel: Channel, is_mobile: bool) -> FT:
-    return Title("tinychat"), Body(cls="font-sans antialiased h-dvh flex bg-background overflow-hidden", hx_ext='ws', ws_connect=f'/ws?mid={m.id}')(
+    return Title("tinychat"), Body(cls="font-sans antialiased h-dvh flex bg-background overflow-hidden", hx_ext='ws', ws_connect='/ws')(
         # sidebar
         Sidebar(m, w, channel, is_mobile) if not is_mobile else None,
         # main content
@@ -781,10 +781,10 @@ def download(fid: str):
     f = file_uploads[fid]
     return FileResponse(f"{settings.file_upload_path}/{f.id}{os.path.splitext(f.original_name)[1]}", filename=f.original_name)
 
-def ws_connect(ws, send):
-    if not ws.session: raise WebSocketException(400, "Missing session")
+def ws_connect(ws, session, send):
+    if not session: raise WebSocketException(400, "Missing session")
     try:
-        m, sid = members[int(ws.query_params.get("mid"))], str(id(ws))
+        m, sid = members[int(session.get("mid"))], str(id(ws))
         connections[sid] = send
         for s in sockets(where=f"mid={m.id}"): sockets.delete(s.sid)
         sockets.insert(Socket(sid=sid, mid=m.id))

locustfile.py

@@ -0,0 +1,51 @@
+# locustfile.py
+# For more options read the following
+#   - https://docs.locust.io/en/stable/writing-a-locustfile.html
+#   - https://docs.locust.io/en/stable/tasksets.html
+
+from locust import HttpUser, task, between
+import gevent, json
+from websockets.sync.client import connect
+
+class TinychatUser(HttpUser):
+    """
+    This class represents simulated users interacting with
+    a website.
+    """
+    # What tasks should be done    
+    # tasks = [TaskSet]
+    # how long between clicks a user should take
+    wait_time = between(2, 20)
+    # The default host of the target client. This can be changed
+    # at any time
+    host = 'http://localhost:5001/'
+
+    def ws_connect(self):
+        session_cookie = self.client.cookies.get_dict().get('session_')
+        self.ws = connect("ws://localhost:5001/ws?mid=3", additional_headers={"Cookie": f"session_={session_cookie}"})
+        self.ws_greenlet = gevent.spawn(self.ws_receive_loop)
+        self.ping_greenlet = gevent.spawn(self.ping_loop)
+
+    def ws_receive_loop(self):
+        while True:
+            message = self.ws.recv()
+            print(f"WS Received: {message}")
+
+    def ping_loop(self):
+        while True:
+            self.ws.send(json.dumps({"cmd": "ping", "d": {"cid": 1}}))
+            gevent.sleep(5)
+
+
+    @task
+    def chat(self):
+        # User goes to the root of the project
+        self.client.get('/')
+        self.client.get('/login')
+        self.client.post('/login', dict(name='locust'))
+        self.client.post('messages/send/1', dict(msg='Hello, world!'))
+
+        self.ws_connect()
+
+        r = self.client.get('/c/1')
+        assert r.status_code == 200