From 6c4ea1507cf10aa2629db0873e11c1bff4fbdd40 Mon Sep 17 00:00:00 2001 From: xainaz Date: Wed, 21 Aug 2024 16:22:05 +0300 Subject: [PATCH 1/3] Removed extra wallet_factoy.py --- aixplain/factories/wallet_factoy.py | 26 -------------------------- 1 file changed, 26 deletions(-) delete mode 100644 aixplain/factories/wallet_factoy.py diff --git a/aixplain/factories/wallet_factoy.py b/aixplain/factories/wallet_factoy.py deleted file mode 100644 index 59ec7c14..00000000 --- a/aixplain/factories/wallet_factoy.py +++ /dev/null @@ -1,26 +0,0 @@ -import aixplain.utils.config as config -from aixplain.modules.wallet import Wallet -from aixplain.utils.file_utils import _request_with_retry -import logging - - -class WalletFactory: - aixplain_key = config.AIXPLAIN_API_KEY - backend_url = config.BACKEND_URL - - @classmethod - def get(cls) -> 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"} - logging.info(f"Start fetching billing information from - {url} - {headers}") - headers = {"Content-Type": "application/json", "x-api-key": config.TEAM_API_KEY} - r = _request_with_retry("get", url, headers=headers) - resp = r.json() - return Wallet(total_balance=resp["totalBalance"], reserved_balance=resp["reservedBalance"]) - except Exception as e: - raise Exception(f"Failed to get the wallet credit information. Error: {str(e)}") From 1950b0a742041ea866cd746acbb742a76829a895 Mon Sep 17 00:00:00 2001 From: xainaz Date: Mon, 26 Aug 2024 16:55:07 +0300 Subject: [PATCH 2/3] Initial Commit --- aixplain/factories/wallet_factory.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/aixplain/factories/wallet_factory.py b/aixplain/factories/wallet_factory.py index 59ec7c14..f968233d 100644 --- a/aixplain/factories/wallet_factory.py +++ b/aixplain/factories/wallet_factory.py @@ -2,6 +2,7 @@ from aixplain.modules.wallet import Wallet from aixplain.utils.file_utils import _request_with_retry import logging +from typing import Text class WalletFactory: @@ -9,18 +10,16 @@ class WalletFactory: 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"]) + return Wallet(total_balance=resp["totalBalance"] or 0, reserved_balance=resp["reservedBalance"] or 0) except Exception as e: raise Exception(f"Failed to get the wallet credit information. Error: {str(e)}") From 843d13ffac03b9c02e3e132480cf2db9c42edb7d Mon Sep 17 00:00:00 2001 From: xainaz Date: Tue, 27 Aug 2024 00:35:37 +0300 Subject: [PATCH 3/3] Added total balance attribute and tested it --- aixplain/factories/wallet_factory.py | 5 ++++- aixplain/modules/wallet.py | 9 +++++---- tests/unit/wallet_test.py | 5 +++-- 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/aixplain/factories/wallet_factory.py b/aixplain/factories/wallet_factory.py index f968233d..b36000f1 100644 --- a/aixplain/factories/wallet_factory.py +++ b/aixplain/factories/wallet_factory.py @@ -20,6 +20,9 @@ def get(cls, api_key: Text = config.TEAM_API_KEY) -> Wallet: 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"] or 0, reserved_balance=resp["reservedBalance"] or 0) + total_balance = float(resp.get("totalBalance", 0.0)) + reserved_balance = float(resp.get("reservedBalance", 0.0)) + + return Wallet(total_balance=total_balance, reserved_balance=reserved_balance) except Exception as e: raise Exception(f"Failed to get the wallet credit information. Error: {str(e)}") diff --git a/aixplain/modules/wallet.py b/aixplain/modules/wallet.py index d7c63524..d61b04ee 100644 --- a/aixplain/modules/wallet.py +++ b/aixplain/modules/wallet.py @@ -24,11 +24,12 @@ class Wallet: def __init__(self, total_balance: float, reserved_balance: float): - """Create a Wallet with the necessary information - + """ Args: - total_balance (float): total credit balance - reserved_balance (float): reserved credit balance + total_balance (float) + reserved_balance (float) + available_balance (float) """ self.total_balance = total_balance self.reserved_balance = reserved_balance + self.available_balance = total_balance-reserved_balance diff --git a/tests/unit/wallet_test.py b/tests/unit/wallet_test.py index 48ee19ab..8e19aab7 100644 --- a/tests/unit/wallet_test.py +++ b/tests/unit/wallet_test.py @@ -12,5 +12,6 @@ def test_wallet_service(): ref_response = {"totalBalance": 5, "reservedBalance": "0"} mock.get(url, headers=headers, json=ref_response) wallet = WalletFactory.get() - assert wallet.total_balance == ref_response["totalBalance"] - assert wallet.reserved_balance == ref_response["reservedBalance"] + assert wallet.total_balance == float(ref_response["totalBalance"]) + assert wallet.reserved_balance == float(ref_response["reservedBalance"]) + assert wallet.available_balance == float(ref_response["totalBalance"])- float(ref_response["reservedBalance"]) \ No newline at end of file