Adjustments to get user credits

aixplain/factories/wallet_factory.py

@@ -2,23 +2,22 @@
 from aixplain.modules.wallet import Wallet
 from aixplain.utils.file_utils import _request_with_retry
 import logging
+from typing import Text
 
 
 class WalletFactory:
     aixplain_key = config.AIXPLAIN_API_KEY
     backend_url = config.BACKEND_URL
 
     @classmethod
-    def get(cls) -> Wallet:
+    def get(cls, api_key: Text = config.TEAM_API_KEY) -> Wallet:     
         """Get wallet information"""
         try:
             resp = None
-            # Check for code 200, other code will be caught when trying to return a Wallet object
             url = f"{cls.backend_url}/sdk/billing/wallet"
-
-            headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"}
+            headers = {"Authorization": f"Token {api_key}", "Content-Type": "application/json"}
             logging.info(f"Start fetching billing information from - {url} - {headers}")
-            headers = {"Content-Type": "application/json", "x-api-key": config.TEAM_API_KEY}
+            headers = {"Content-Type": "application/json", "x-api-key": api_key}
             r = _request_with_retry("get", url, headers=headers)
             resp = r.json()
             return Wallet(total_balance=resp["totalBalance"], reserved_balance=resp["reservedBalance"])

tests/unit/wallet_test.py

@@ -11,6 +11,6 @@ def test_wallet_service():
         headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"}
         ref_response = {"totalBalance": 5, "reservedBalance": "0"}
         mock.get(url, headers=headers, json=ref_response)
-        wallet = WalletFactory.get()
+        wallet = WalletFactory.get(config.AIXPLAIN_API_KEY)
     assert wallet.total_balance == ref_response["totalBalance"]
     assert wallet.reserved_balance == ref_response["reservedBalance"]