From 5bb325a6970084eb4d9587aa8728cc937258444a Mon Sep 17 00:00:00 2001
From: Joe Gough <36932486+honourfish@users.noreply.github.com>
Date: Fri, 19 Jul 2024 13:28:11 +0100
Subject: [PATCH] Better error handling and logging for get operation status
 (#19)

* Better error handling and logging for get operation status

---------

Co-authored-by: jgough <joe.gough@datatrails.ai>
---
 requirements-dev.txt            |  2 +-
 scitt/check_operation_status.py | 49 ++++++++++++++++++++++++++-------
 unittests/__init__.py           |  1 +
 3 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/requirements-dev.txt b/requirements-dev.txt
index 1ae4778..741cb88 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -2,7 +2,7 @@
 
 # code quality
 autopep8~=2.0
-black~=23.9
+black~=24.4.2
 pycodestyle~=2.10
 pylint~=3.0
 pyright~=1.1
diff --git a/scitt/check_operation_status.py b/scitt/check_operation_status.py
index 267cc38..e74e62b 100755
--- a/scitt/check_operation_status.py
+++ b/scitt/check_operation_status.py
@@ -2,6 +2,8 @@
 
 import os
 import argparse
+import logging
+import sys
 
 from time import sleep as time_sleep
 
@@ -10,7 +12,7 @@
 
 # all timeouts and durations are in seconds
 REQUEST_TIMEOUT = 30
-POLL_TIMEOUT = 360
+POLL_TIMEOUT = 60
 POLL_INTERVAL = 10
 
 
@@ -41,24 +43,36 @@ def get_operation_status(operation_id: str, headers: dict) -> dict:
     return response.json()
 
 
-def poll_operation_status(operation_id: str, headers: dict) -> str:
+def poll_operation_status(
+    operation_id: str, headers: dict, logger: logging.Logger
+) -> str:
     """
     polls for the operation status to be 'succeeded'.
     """
 
     poll_attempts: int = int(POLL_TIMEOUT / POLL_INTERVAL)
 
+    logger.info("starting to poll for operation status 'succeeded'")
+
     for _ in range(poll_attempts):
-        operation_status = get_operation_status(operation_id, headers)
 
-        # pylint: disable=fixme
-        # TODO: ensure get_operation_status handles error cases from the rest request
-        if "status" in operation_status and operation_status["status"] == "succeeded":
-            return operation_status["entryID"]
+        try:
+            operation_status = get_operation_status(operation_id, headers)
+
+            # pylint: disable=fixme
+            # TODO: ensure get_operation_status handles error cases from the rest request
+            if (
+                "status" in operation_status
+                and operation_status["status"] == "succeeded"
+            ):
+                return operation_status["entryID"]
+
+        except requests.HTTPError as e:
+            logger.debug("failed getting operation status, error: %s", e)
 
         time_sleep(POLL_INTERVAL)
 
-    raise TimeoutError("signed statement not registered within polling duration.")
+    raise TimeoutError("signed statement not registered within polling duration")
 
 
 def main():
@@ -91,12 +105,27 @@ def main():
         default=default_token_file_name,
     )
 
+    # log level
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        help="log level. for any individual poll errors use DEBUG, defaults to WARNING",
+        default="WARNING",
+    )
+
     args = parser.parse_args()
 
+    logger = logging.getLogger("check operation status")
+    logging.basicConfig(level=logging.getLevelName(args.log_level))
+
     headers = get_token_from_file(args.token_file_name)
 
-    entry_id = poll_operation_status(args.operation_id, headers)
-    print(entry_id)
+    try:
+        entry_id = poll_operation_status(args.operation_id, headers, logger)
+        print(entry_id)
+    except TimeoutError as e:
+        print(e, file=sys.stderr)
+        sys.exit(1)
 
 
 if __name__ == "__main__":
diff --git a/unittests/__init__.py b/unittests/__init__.py
index 782379d..1ada2a3 100644
--- a/unittests/__init__.py
+++ b/unittests/__init__.py
@@ -1,6 +1,7 @@
 """
 Unit tests
 """
+
 import unittest
 
 # Hides Docstring