fix: balance misconfigured request

viewdns/tests/test_balance.py

@@ -17,7 +17,7 @@ def test_balace(self):
 
         data = self.load_from_file('balance.json')
 
-        url = self.base_url
+        url = self.base_url + 'account'
         responses.add(responses.GET, url, body=data, status=200)
 
         balance = self.client.balance()

viewdns/viewdns/Balance.py

@@ -3,9 +3,17 @@
 class Balance(BaseObject):
 
     def __init__(self, *args, **kwargs):
-
-        self.limit = int(kwargs['monthly']['limit'])
-        self.usage = int(kwargs['monthly']['usage'])
+        if 'monthly' in kwargs:
+            # this is paid account
+            self.limit = int(kwargs['monthly']['limit'])
+            self.usage = int(kwargs['monthly']['usage'])
+        elif 'trial' in kwargs:
+            # this is a free/trial account
+            self.limit = int(kwargs['trial']['limit'])
+            self.usage = int(kwargs['trial']['usage']) 
+        else:
+            existing_keys = ",".join(list(kwargs.keys()))
+            raise KeyError(f"Missing trial or prepaid key in response. Received Keys are {existing_keys}")
         self.balance = int(kwargs['prepaid']['balance'])
         # warning is set at 80% limit
         self.warning = self.usage >= self.limit * 0.8

viewdns/viewdns/Client.py

@@ -1,10 +1,12 @@
 import requests
 
 try:
-    import urlparse  # type: ignore [import-not-found]
+    import urlparse  # type: ignore [import-not-found, import]
 except ImportError:
     from urllib import parse as urlparse
 
+from json.decoder import JSONDecodeError
+
 from .DNSRecord import DNSRecord
 from .HTTPHeader import HTTPHeader
 from .IPLocation import IPLocation
@@ -338,7 +340,7 @@ def balance(self):
         params = dict()
         params['action'] = 'balance'
 
-        res = self._execute(None, params=params)
+        res = self._execute('account', params=params)
 
         balance = Balance(**res['response'])
 
@@ -362,7 +364,10 @@ def _execute(self, url = None, params=None):
 
         try:
             data = req.json()
-        except ValueError as e:
-            raise Exception(e)
-
-        return data
+        except JSONDecodeError as exc:
+            print(req.content)
+            raise Exception("Error while decoding json") from exc
+        except ValueError as exc:
+            raise Exception(exc)
+        else:
+            return data