diff --git a/bittensor/utils/balance.py b/bittensor/utils/balance.py index 0bba07622d..a52913c37d 100644 --- a/bittensor/utils/balance.py +++ b/bittensor/utils/balance.py @@ -68,6 +68,9 @@ def __repr__(self): return self.__str__() def __eq__(self, other: Union[int, float, "Balance"]): + if other is None: + return False + if hasattr(other, "rao"): return self.rao == other.rao else: diff --git a/tests/unit_tests/bittensor_tests/test_balance.py b/tests/unit_tests/bittensor_tests/test_balance.py index 60f61fac67..8a52d117ab 100644 --- a/tests/unit_tests/bittensor_tests/test_balance.py +++ b/tests/unit_tests/bittensor_tests/test_balance.py @@ -327,3 +327,12 @@ def test_balance_rfloordiv_other_not_balance(self, balance: Union[int, float], b assert isinstance(quot_, Balance) assert CLOSE_IN_VALUE(quot_.rao, 5) == rao2_ // rao_ + @given(balance=valid_tao_numbers_strategy) + def test_balance_not_eq_none(self, balance: Union[int, float]): + balance_ = Balance(balance) + assert not balance_ == None + + @given(balance=valid_tao_numbers_strategy) + def test_balance_neq_none(self, balance: Union[int, float]): + balance_ = Balance(balance) + assert balance_ != None