From c79c3633884d38c878657c738eee238220b81793 Mon Sep 17 00:00:00 2001
From: Kameron Kenny <1267885+kkenny@users.noreply.github.com>
Date: Sun, 20 Oct 2024 21:59:54 -0400
Subject: [PATCH] add try/except logic

---
 power_client.py | 55 +++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 46 insertions(+), 9 deletions(-)

diff --git a/power_client.py b/power_client.py
index ba9d44a..b6662ac 100644
--- a/power_client.py
+++ b/power_client.py
@@ -3,6 +3,9 @@ import time
 
 import requests
 import socket
+
+from urllib3.exceptions import NewConnectionError
+
 from killswitch import KillSwitch
 from timber import LumberJack, axe
 from config import _CLIENT_CONFIG as _c
@@ -33,8 +36,16 @@ class PowerClient:
         self.update_config()
 
     def get_config(self):
-        c = requests.get(self.config_url)
-        if c.status_code == 200:
+        try:
+            c = requests.get(self.config_url)
+        except NewConnectionError as NCE:
+            axe.error(f'Error: {NCE}')
+            c = None
+        except requests.exceptions.ConnectionError as CE:
+            axe.error(f'Error: {CE}')
+            c = None
+
+        if c and c.status_code == 200:
             return c.json()
             # return json.loads(c.json())
         else:
@@ -115,7 +126,18 @@ class PowerClient:
         return hn
 
     def get_status_code(self):
-        return requests.get(self.url)
+        try:
+            r = requests.get(self.url)
+        except NewConnectionError as NCE:
+            axe.error(f'Error: {NCE}')
+            r = None
+        except requests.exceptions.ConnectionError as CE:
+            axe.error(f'Error: {CE}')
+            r = None
+        return r
+
+    def set_unhealthy(self):
+       pass
 
     def check_power(self):
         axe.debug(f'Checking {self.url} for power with {self.RETRY_ATTEMPTS} retry attempts, DRYRUN={self.DRYRUN}')
@@ -130,13 +152,28 @@ class PowerClient:
                 self.update_config()
                 self.config_check_count = 0
 
-            if response.ok and response.status_code == 200:
-                if self.RETRY_COUNT == 1:
-                    axe.info(f'Power is up')
+            if response:
+                if response.ok and response.status_code == 200:
+                    if self.RETRY_COUNT == 1:
+                        axe.info(f'Power is up')
+                    else:
+                        axe.info(f'Power is back up after {self.RETRY_COUNT} checks')
+                    self.RETRY_COUNT = 0
+                    self.RETRY_SECONDS = self.HEALTHY_RETRY_SECONDS
                 else:
-                    axe.info(f'Power is back up after {self.RETRY_COUNT} checks')
-                self.RETRY_COUNT = 0
-                self.RETRY_SECONDS = self.HEALTHY_RETRY_SECONDS
+                    self.RETRY_SECONDS = self.UNHEALTHY_RETRY_SECONDS
+                    if self.RETRY_COUNT < int(math.floor(self.RETRY_ATTEMPTS / 2)):
+                        axe.warning(f'Attempt {self.RETRY_COUNT} request failed, power may be down. \
+                                 Will try {remaining_tries} more times before initiating shutdown.')
+                    elif self.RETRY_COUNT < self.RETRY_ATTEMPTS:
+                        axe.error(f'Attempt {self.RETRY_COUNT} requests failed, power appears down.\
+                                 Will try {remaining_tries} more times before initiating shutdown.')
+                    else:
+                        axe.critical(f'Attempt {self.RETRY_COUNT} requests failed, power appears down.\
+                                     Engaging Killswitch!')
+                        ks = KillSwitch(DRYRUN=self.DRYRUN)
+                        ks.power_down()
+                        break
             else:
                 self.RETRY_SECONDS = self.UNHEALTHY_RETRY_SECONDS
                 if self.RETRY_COUNT < int(math.floor(self.RETRY_ATTEMPTS / 2)):