Add log message to checker exceptions and assert_* helpers (#39)

* Add log message to checker exceptions and assert_* helpers * add return type annotation to caller_loc --------- Co-authored-by: ldruschk <[email protected]> Co-authored-by: Lucas Druschke <[email protected]>
enowars · Oct 15, 2023 · bdd7087 · bdd7087
1 parent 453917c
commit bdd7087
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 9 deletions.
diff --git a/enochecker3/enochecker.py b/enochecker3/enochecker.py
@@ -332,7 +332,7 @@ async def _call_method(self, task: BaseCheckerTaskMessage) -> Optional[str]:
             trace = traceback.format_exc()
             logger = self._get_logger_adapter(task)
             logger.info(f"Checker internal error\n{trace}")
-            raise InternalErrorException("Checker internal error", e)
+            raise InternalErrorException("Checker internal error", inner=e)
 
     async def _call_putflag(
         self, task: PutflagCheckerTaskMessage
@@ -559,18 +559,24 @@ async def checker(task: CheckerTaskMessage) -> CheckerResultMessage:
                     )
             except MumbleException as e:
                 trace = traceback.format_exc()
+                if e.log_message:
+                    logger.info(e.log_message)
                 logger.info(f"Encountered mumble exception:\n{trace}")
                 return CheckerResultMessage(
                     result=CheckerTaskResult.MUMBLE, message=e.message
                 )
             except OfflineException as e:
                 trace = traceback.format_exc()
+                if e.log_message:
+                    logger.info(e.log_message)
                 logger.info(f"Encountered offline exception:\n{trace}")
                 return CheckerResultMessage(
                     result=CheckerTaskResult.OFFLINE, message=e.message
                 )
             except InternalErrorException as e:
                 trace = traceback.format_exc()
+                if e.log_message:
+                    logger.info(e.log_message)
                 logger.info(f"Encountered internal error exception:\n{trace}")
                 return CheckerResultMessage(
                     result=CheckerTaskResult.INTERNAL_ERROR, message=e.message

diff --git a/enochecker3/types.py b/enochecker3/types.py
@@ -106,7 +106,7 @@ class CheckerInfoMessage(BaseModel):
 
 class CheckerResultMessage(BaseModel):
     result: CheckerTaskResult
-    message: Optional[str]
+    message: Optional[str] = None
     attack_info: Optional[str] = None
     flag: Optional[str] = None
 
@@ -130,8 +130,9 @@ class CheckerTaskMessage(BaseModel):
 
 
 class BaseException(Exception):
-    def __init__(self, message: Optional[str]):
+    def __init__(self, message: Optional[str], log_message: Optional[str] = None):
         self.message: Optional[str] = message
+        self.log_message: Optional[str] = log_message
 
 
 class MumbleException(BaseException):
@@ -143,8 +144,13 @@ class OfflineException(BaseException):
 
 
 class InternalErrorException(BaseException):
-    def __init__(self, message: Optional[str], inner: Optional[Exception] = None):
-        super().__init__(message)
+    def __init__(
+        self,
+        message: Optional[str],
+        log_message: Optional[str] = None,
+        inner: Optional[Exception] = None,
+    ):
+        super().__init__(message, log_message)
         self.inner: Optional[Exception] = inner
 
 

diff --git a/enochecker3/utils.py b/enochecker3/utils.py
@@ -1,20 +1,33 @@
 import hashlib
 import re
+import traceback
 from typing import Any, Optional, Union
 
 from .types import MumbleException
 
 
+def caller_loc() -> str:
+    trace = traceback.StackSummary.extract(traceback.walk_stack(None))
+    caller = list(trace)[2]
+    return f"{caller.filename}:{caller.lineno}"
+
+
 def assert_in(o1: Any, o2: Any, message: Optional[str] = None) -> None:
     if not o2 or o1 not in o2:
         raise MumbleException(
             message or "Checker assertion failed",
+            log_message=f"Assertion (assert_in) failed! ({caller_loc()})"
+            + f"\n  Needle: ({type(o1)}) {o1}\n  Haystack: ({type(o2)}) {o2}",
         )
 
 
 def assert_equals(o1: Any, o2: Any, message: Optional[str] = None) -> None:
     if o1 != o2:
-        raise MumbleException(message or "Checker assertion failed")
+        raise MumbleException(
+            message or "Checker assertion failed",
+            log_message=f"Assertion (assert_equals) failed! ({caller_loc()})"
+            + f"\n  Left: ({type(o1)}) {o1}\n  Right: ({type(o2)}) {o2}",
+        )
 
 
 class FlagSearcher:

diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
 httpx>=0.23.0,<0.24.0
-motor[asyncio]>=3.1.2,<4.0.0
+motor>=3.1.2,<4.0.0
 fastapi>=0.75.2,<0.76.0
 uvicorn>=0.17.6,<0.18.0
 gunicorn==20.1.0
diff --git a/tests/test_chaindb.py b/tests/test_chaindb.py
@@ -2,13 +2,14 @@
 
 import pymongo
 import pytest
+import pytest_asyncio
 from motor.core import AgnosticClient, AgnosticCollection, AgnosticDatabase
 from motor.motor_asyncio import AsyncIOMotorClient
 
 from enochecker3 import ChainDB
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def collection():
     mongo: AgnosticClient = AsyncIOMotorClient()
     db: AgnosticDatabase = mongo["enochecker3_tests"]
@@ -23,7 +24,7 @@ async def collection():
     return chain_collection
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def chaindb(collection):
     task_chain_id = secrets.token_hex(8)
     return ChainDB(collection, task_chain_id)