Implement Balance (#321)

* Added Balance

* Import Balance into __init__, Manager.py, add tests

Co-authored-by: Work Slave <287328@auschwitz.camp>

Author: Teran McKinney

digitalocean/Balance.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+from .baseapi import BaseAPI
+
+
+class Balance(BaseAPI):
+    def __init__(self, *args, **kwargs):
+        self.month_to_date_balance = None
+        self.account_balance = None
+        self.month_to_date_usage = None
+        self.generated_at = None
+
+        super(Balance, self).__init__(*args, **kwargs)
+
+    @classmethod
+    def get_object(cls, api_token):
+        """
+            Class method that will return an Balance object.
+        """
+        acct = cls(token=api_token)
+        acct.load()
+        return acct
+
+    def load(self):
+        # URL https://api.digitalocean.com/customers/my/balance
+        balance = self.get_data("customers/my/balance")
+
+        for attr in balance.keys():
+            setattr(self, attr, balance[attr])
+
+    def __str__(self):
+        return "<Balance: %s>" % (self.account_balance)

digitalocean/Manager.py

@@ -7,6 +7,7 @@
 from .baseapi import BaseAPI
 from .Account import Account
 from .Action import Action
+from .Balance import Balance
 from .Certificate import Certificate
 from .Domain import Domain
 from .Droplet import Droplet
@@ -34,6 +35,12 @@ def get_account(self):
         """
         return Account.get_object(api_token=self.token)
 
+    def get_balance(self):
+        """
+            Returns a Balance object.
+        """
+        return Balance.get_object(api_token=self.token)
+
     def get_all_regions(self):
         """
             This function returns a list of Region object.

digitalocean/__init__.py

@@ -14,6 +14,7 @@
 from .Image import Image
 from .Action import Action
 from .Account import Account
+from .Balance import Balance
 from .Domain import Domain
 from .Record import Record
 from .SSHKey import SSHKey

digitalocean/tests/data/balance/balance.json

@@ -0,0 +1,6 @@
+{
+  "month_to_date_balance": "23.44",
+  "account_balance": "12.23",
+  "month_to_date_usage": "11.21",
+  "generated_at": "2019-07-09T15:01:12Z"
+}

digitalocean/tests/test_manager.py

@@ -34,6 +34,25 @@ def test_get_account(self):
         self.assertEqual(acct.email_verified, True)
         self.assertEqual(acct.status, "active")
 
+    @responses.activate
+    def test_get_balance(self):
+        data = self.load_from_file('balance/balance.json')
+
+        url = self.base_url + 'customers/my/balance'
+        responses.add(responses.GET, url,
+                      body=data,
+                      status=200,
+                      content_type='application/json')
+
+        balance = self.manager.get_balance()
+
+        self.assert_get_url_equal(responses.calls[0].request.url, url)
+        self.assertEqual(balance.token, balance.token)
+        self.assertEqual(balance.month_to_date_balance, '23.44')
+        self.assertEqual(balance.account_balance, '12.23')
+        self.assertEqual(balance.month_to_date_usage, '11.21')
+        self.assertEqual(balance.generated_at, '2019-07-09T15:01:12Z')
+
     @responses.activate
     def test_auth_fail(self):
         data = self.load_from_file('errors/unauthorized.json')