From 0c0a42a4c91e69252523c854197d9a75298b9e29 Mon Sep 17 00:00:00 2001
From: Guilhem Saurel <guilhem.saurel@laas.fr>
Date: Wed, 14 Jul 2021 23:25:24 +0200
Subject: [PATCH] improve error management and tests

---
 matrix_webhook.py | 29 ++++++++++++++++++++++-------
 tests/tests.py    | 20 +++++++++++++++++---
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/matrix_webhook.py b/matrix_webhook.py
index c52fe43..70d3234 100755
--- a/matrix_webhook.py
+++ b/matrix_webhook.py
@@ -17,6 +17,7 @@ from aiohttp import web
 from markdown import markdown
 from nio import AsyncClient
 from nio.exceptions import LocalProtocolError
+from nio.responses import RoomSendError
 
 parser = argparse.ArgumentParser(description=__doc__)
 parser.add_argument(
@@ -81,6 +82,7 @@ MATRIX_PW = args.matrix_pw
 API_KEY = args.api_key
 CLIENT = AsyncClient(args.matrix_url, args.matrix_id)
 LOGGER = logging.getLogger("matrix-webhook")
+ERROR_MAP = {"M_FORBIDDEN": HTTPStatus.FORBIDDEN}
 
 
 async def handler(request):
@@ -112,13 +114,26 @@ async def handler(request):
         "format": "org.matrix.custom.html",
         "formatted_body": markdown(str(data["text"]), extensions=["extra"]),
     }
-    try:
-        await send_room_message(room_id, content)
-    except LocalProtocolError as e:  # Connection lost, try another login
-        LOGGER.error(f"Send error: {e}")
-        LOGGER.warning("Reconnecting and trying again")
-        await CLIENT.login(MATRIX_PW)
-        await send_room_message(room_id, content)
+    for _ in range(10):
+        try:
+            resp = await send_room_message(room_id, content)
+            if isinstance(resp, RoomSendError):
+                if resp.status_code == "M_UNKNOWN_TOKEN":
+                    LOGGER.warning("Reconnecting")
+                    await CLIENT.login(MATRIX_PW)
+                else:
+                    return create_json_response(
+                        ERROR_MAP[resp.status_code], resp.message
+                    )
+            else:
+                break
+        except LocalProtocolError as e:
+            LOGGER.error(f"Send error: {e}")
+        LOGGER.warning("Trying again")
+    else:
+        return create_json_response(
+            HTTPStatus.GATEWAY_TIMEOUT, "Homeserver not responding"
+        )
 
     return create_json_response(HTTPStatus.OK, "OK")
 
diff --git a/tests/tests.py b/tests/tests.py
index f6ec6f9..c92fb95 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -20,9 +20,11 @@ class BotTest(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(
             bot_req({"text": 3, "key": None}), {"status": 401, "ret": "Invalid API key"}
         )
-
-        # TODO: we are not sending to a real room, so this should not be "OK"
-        self.assertEqual(bot_req({"text": 3}, KEY), {"status": 200, "ret": "OK"})
+        # TODO: if the client from matrix_webhook has olm support, this won't be a 403 from synapse,
+        # but a LocalProtocolError from matrix_webhook
+        self.assertEqual(
+            bot_req({"text": 3}, KEY), {"status": 403, "ret": "Unknown room"}
+        )
 
     async def test_message(self):
         """Send a markdown message, and check the result."""
@@ -45,3 +47,15 @@ class BotTest(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(message.sender, FULL_ID)
         self.assertEqual(message.body, text)
         self.assertEqual(message.formatted_body, "<h1>Hello</h1>")
+
+    async def test_reconnect(self):
+        """Check the reconnecting path."""
+        client = nio.AsyncClient(MATRIX_URL, MATRIX_ID)
+        await client.login(MATRIX_PW)
+        room = await client.room_create()
+        await client.logout(all_devices=True)
+        await client.close()
+        self.assertEqual(
+            bot_req({"text": "Re"}, KEY, room.room_id),
+            {"status": 200, "ret": "OK"},
+        )