Merge pull request #20 from unniznd/feature_get_balance

feature: getBalance API added and tested successfully

Author: aakash-taneja

src/lighthouseweb3/__init__.py

@@ -7,7 +7,8 @@
     deal_status, 
     get_uploads as getUploads, 
     download as _download,
-    get_file_info as getFileInfo
+    get_file_info as getFileInfo,
+    get_balance as getBalance
 )
 
 
@@ -44,6 +45,17 @@ def uploadBlob(self, source: io.BufferedReader, filename: str, tag: str = ''):
             return d.uploadBlob(source, filename, self.token, tag)
         except Exception as e:
             raise e
+    
+    def getBalance(self):
+        """
+        Retrieve the balance information of a user from the Lighthouse.
+        :param publicKey: str, The public key of the user.
+        :return: dict[str, any], A dictionary containing the data usage and data limit details.
+        """
+        try:
+            return getBalance.get_balance(self.token)
+        except Exception as e:
+            raise e
 
     @staticmethod
     def downloadBlob(dist: io.BufferedWriter, cid: str, chunk_size=1024*1024*10):

src/lighthouseweb3/functions/get_balance.py

@@ -0,0 +1,14 @@
+from .config import Config
+import requests as req
+
+def get_balance(token:str):
+  headers = {
+        "Authorization": f"Bearer {token}",
+    }
+  url = f"{Config.lighthouse_api}/api/user/user_data_usage"
+  try:
+    response = req.get(url, headers=headers)
+  except Exception as e:
+    raise Exception("Failed to get account balance")
+
+  return response.json()

tests/test_get_balance.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python3
+import os
+import unittest
+from src.lighthouseweb3 import Lighthouse
+from .setup import parse_env
+
+
+class TestGetBalance(unittest.TestCase):
+
+    def test_get_balance(self):
+        """test get_balance function"""
+        parse_env()
+        l = Lighthouse(os.environ.get("LIGHTHOUSE_TOKEN"))
+        res = l.getBalance()
+        self.assertIsInstance(res, dict, "data is a dict")
+        self.assertIsInstance(res.get("dataLimit"), int, "data limit is a integer")
+        self.assertIsInstance(res.get("dataUsed"), int, "data used is a integer")
+    
+    def test_get_balance_invalid_token(self):
+        """test get_balance function with invalid token"""
+        parse_env()
+        l = Lighthouse('invalid_token')
+        res = l.getBalance()
+        self.assertIsInstance(res, dict, "res is a dict")
+        self.assertIn("authentication failed", str(res).lower())
+