From 01bd63773078dad39c717b68ce7823c1590b4662 Mon Sep 17 00:00:00 2001
From: norohind <60548839+norohind@users.noreply.github.com>
Date: Fri, 13 Aug 2021 21:47:47 +0300
Subject: [PATCH] Add CL arg `--capi-pretend-down` in order to pretend that
 CAPI is down

---
 EDMarketConnector.py | 10 ++++++++++
 companion.py         | 10 ++++++++++
 config.py            |  1 +
 3 files changed, 21 insertions(+)

diff --git a/EDMarketConnector.py b/EDMarketConnector.py
index b2d4cba8..886bf1a7 100755
--- a/EDMarketConnector.py
+++ b/EDMarketConnector.py
@@ -127,8 +127,18 @@ if __name__ == '__main__':  # noqa: C901
                         nargs='*'
                         )
 
+    parser.add_argument(
+        '--capi-pretend-down',
+        help='Force to raise ServerError on any CAPI query',
+        action='store_true')
+
     args = parser.parse_args()
 
+    if args.capi_pretend_down:
+        import config as conf_module
+        logger.info('Pretending CAPI is down')
+        conf_module.capi_pretend_down = True
+
     level_to_set: Optional[int] = None
     if args.trace or args.trace_on:
         level_to_set = logging.TRACE  # type: ignore # it exists
diff --git a/companion.py b/companion.py
index f2f571b6..7f55fc7f 100644
--- a/companion.py
+++ b/companion.py
@@ -25,6 +25,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, OrderedDic
 import requests
 
 from config import appname, appversion, config
+import config as conf_module
 from edmc_data import companion_category_map as category_map
 from EDMCLogging import get_main_logger
 from monitor import monitor
@@ -276,6 +277,9 @@ class Auth(object):
 
             logger.debug('Attempting refresh with Frontier...')
             try:
+                if conf_module.capi_pretend_down:
+                    raise ServerError(_('Pretending CAPI is down'))
+
                 r = self.session.post(SERVER_AUTH + URL_TOKEN, data=data, timeout=auth_timeout)
                 if r.status_code == requests.codes.ok:
                     data = r.json()
@@ -360,6 +364,9 @@ class Auth(object):
             # requests_log.setLevel(logging.DEBUG)
             # requests_log.propagate = True
 
+            if conf_module.capi_pretend_down:
+                raise ServerError(_('Pretending CAPI is down'))
+
             r = self.session.post(SERVER_AUTH + URL_TOKEN, data=request_data, timeout=auth_timeout)
             data_token = r.json()
             if r.status_code == requests.codes.ok:
@@ -565,6 +572,9 @@ class Session(object):
 
         try:
             logger.trace('Trying...')
+            if conf_module.capi_pretend_down:
+                raise ServerError('Pretending CAPI is down')
+
             r = self.session.get(self.server + endpoint, timeout=timeout)  # type: ignore
 
         except requests.ConnectionError as e:
diff --git a/config.py b/config.py
index 012eb84d..d1cfd403 100644
--- a/config.py
+++ b/config.py
@@ -45,6 +45,7 @@ debug_senders: List[str] = []
 # *all* if only interested in some things.
 trace_on: List[str] = []
 
+capi_pretend_down: bool = False
 # This must be done here in order to avoid an import cycle with EDMCLogging.
 # Other code should use EDMCLogging.get_main_logger
 if os.getenv("EDMC_NO_UI"):