From d842fc288fce35f90222528682e783ec5cba0597 Mon Sep 17 00:00:00 2001 From: Tom Brien Date: Fri, 6 Aug 2021 17:34:21 +0100 Subject: [PATCH] Ignore Coinbase vault wallets (#54133) * Exclude vault balances * Update option flow validation * Update test name * Add missed check * Fix dangerous default --- .../components/coinbase/config_flow.py | 8 +++- homeassistant/components/coinbase/const.py | 2 + homeassistant/components/coinbase/sensor.py | 16 ++++++-- tests/components/coinbase/common.py | 10 ++--- tests/components/coinbase/const.py | 14 ++++++- tests/components/coinbase/test_config_flow.py | 4 +- tests/components/coinbase/test_init.py | 41 +++++++++++++++---- 7 files changed, 74 insertions(+), 21 deletions(-) diff --git a/homeassistant/components/coinbase/config_flow.py b/homeassistant/components/coinbase/config_flow.py index 4ea36dad266..5901aeeed9a 100644 --- a/homeassistant/components/coinbase/config_flow.py +++ b/homeassistant/components/coinbase/config_flow.py @@ -14,6 +14,8 @@ from . import get_accounts from .const import ( API_ACCOUNT_CURRENCY, API_RATES, + API_RESOURCE_TYPE, + API_TYPE_VAULT, CONF_CURRENCIES, CONF_EXCHANGE_BASE, CONF_EXCHANGE_RATES, @@ -65,7 +67,11 @@ async def validate_options( accounts = await hass.async_add_executor_job(get_accounts, client) - accounts_currencies = [account[API_ACCOUNT_CURRENCY] for account in accounts] + accounts_currencies = [ + account[API_ACCOUNT_CURRENCY] + for account in accounts + if account[API_RESOURCE_TYPE] != API_TYPE_VAULT + ] available_rates = await hass.async_add_executor_job(client.get_exchange_rates) if CONF_CURRENCIES in options: for currency in options[CONF_CURRENCIES]: diff --git a/homeassistant/components/coinbase/const.py b/homeassistant/components/coinbase/const.py index a7ed0b15986..dc2922d1531 100644 --- a/homeassistant/components/coinbase/const.py +++ b/homeassistant/components/coinbase/const.py @@ -18,6 +18,8 @@ API_ACCOUNT_NATIVE_BALANCE = "native_balance" API_ACCOUNT_NAME = "name" API_ACCOUNTS_DATA = "data" API_RATES = "rates" +API_RESOURCE_TYPE = "type" +API_TYPE_VAULT = "vault" WALLETS = { "1INCH": "1INCH", diff --git a/homeassistant/components/coinbase/sensor.py b/homeassistant/components/coinbase/sensor.py index c86f21bac1d..f836a604f6a 100644 --- a/homeassistant/components/coinbase/sensor.py +++ b/homeassistant/components/coinbase/sensor.py @@ -12,6 +12,8 @@ from .const import ( API_ACCOUNT_NAME, API_ACCOUNT_NATIVE_BALANCE, API_RATES, + API_RESOURCE_TYPE, + API_TYPE_VAULT, CONF_CURRENCIES, CONF_EXCHANGE_RATES, DOMAIN, @@ -41,7 +43,9 @@ async def async_setup_entry(hass, config_entry, async_add_entities): entities = [] provided_currencies = [ - account[API_ACCOUNT_CURRENCY] for account in instance.accounts + account[API_ACCOUNT_CURRENCY] + for account in instance.accounts + if account[API_RESOURCE_TYPE] != API_TYPE_VAULT ] desired_currencies = [] @@ -82,7 +86,10 @@ class AccountSensor(SensorEntity): self._coinbase_data = coinbase_data self._currency = currency for account in coinbase_data.accounts: - if account[API_ACCOUNT_CURRENCY] == currency: + if ( + account[API_ACCOUNT_CURRENCY] == currency + and account[API_RESOURCE_TYPE] != API_TYPE_VAULT + ): self._name = f"Coinbase {account[API_ACCOUNT_NAME]}" self._id = ( f"coinbase-{account[API_ACCOUNT_ID]}-wallet-" @@ -135,7 +142,10 @@ class AccountSensor(SensorEntity): """Get the latest state of the sensor.""" self._coinbase_data.update() for account in self._coinbase_data.accounts: - if account[API_ACCOUNT_CURRENCY] == self._currency: + if ( + account[API_ACCOUNT_CURRENCY] == self._currency + and account[API_RESOURCE_TYPE] != API_TYPE_VAULT + ): self._state = account[API_ACCOUNT_BALANCE][API_ACCOUNT_AMOUNT] self._native_balance = account[API_ACCOUNT_NATIVE_BALANCE][ API_ACCOUNT_AMOUNT diff --git a/tests/components/coinbase/common.py b/tests/components/coinbase/common.py index 5fcab6605bd..231a5128585 100644 --- a/tests/components/coinbase/common.py +++ b/tests/components/coinbase/common.py @@ -6,7 +6,7 @@ from homeassistant.components.coinbase.const import ( ) from homeassistant.const import CONF_API_KEY, CONF_API_TOKEN -from .const import GOOD_EXCHNAGE_RATE, GOOD_EXCHNAGE_RATE_2, MOCK_ACCOUNTS_RESPONSE +from .const import GOOD_EXCHANGE_RATE, GOOD_EXCHANGE_RATE_2, MOCK_ACCOUNTS_RESPONSE from tests.common import MockConfigEntry @@ -60,11 +60,11 @@ def mock_get_exchange_rates(): """Return a heavily reduced mock list of exchange rates for testing.""" return { "currency": "USD", - "rates": {GOOD_EXCHNAGE_RATE_2: "0.109", GOOD_EXCHNAGE_RATE: "0.00002"}, + "rates": {GOOD_EXCHANGE_RATE_2: "0.109", GOOD_EXCHANGE_RATE: "0.00002"}, } -async def init_mock_coinbase(hass): +async def init_mock_coinbase(hass, currencies=None, rates=None): """Init Coinbase integration for testing.""" config_entry = MockConfigEntry( domain=DOMAIN, @@ -72,8 +72,8 @@ async def init_mock_coinbase(hass): title="Test User", data={CONF_API_KEY: "123456", CONF_API_TOKEN: "AbCDeF"}, options={ - CONF_CURRENCIES: [], - CONF_EXCHANGE_RATES: [], + CONF_CURRENCIES: currencies or [], + CONF_EXCHANGE_RATES: rates or [], }, ) config_entry.add_to_hass(hass) diff --git a/tests/components/coinbase/const.py b/tests/components/coinbase/const.py index 7d36d0be9a7..082c986aa59 100644 --- a/tests/components/coinbase/const.py +++ b/tests/components/coinbase/const.py @@ -3,8 +3,8 @@ GOOD_CURRENCY = "BTC" GOOD_CURRENCY_2 = "USD" GOOD_CURRENCY_3 = "EUR" -GOOD_EXCHNAGE_RATE = "BTC" -GOOD_EXCHNAGE_RATE_2 = "ATOM" +GOOD_EXCHANGE_RATE = "BTC" +GOOD_EXCHANGE_RATE_2 = "ATOM" BAD_CURRENCY = "ETH" BAD_EXCHANGE_RATE = "ETH" @@ -15,6 +15,15 @@ MOCK_ACCOUNTS_RESPONSE = [ "id": "123456789", "name": "BTC Wallet", "native_balance": {"amount": "100.12", "currency": GOOD_CURRENCY_2}, + "type": "wallet", + }, + { + "balance": {"amount": "100.00", "currency": GOOD_CURRENCY}, + "currency": GOOD_CURRENCY, + "id": "abcdefg", + "name": "BTC Vault", + "native_balance": {"amount": "100.12", "currency": GOOD_CURRENCY_2}, + "type": "vault", }, { "balance": {"amount": "9.90", "currency": GOOD_CURRENCY_2}, @@ -22,5 +31,6 @@ MOCK_ACCOUNTS_RESPONSE = [ "id": "987654321", "name": "USD Wallet", "native_balance": {"amount": "9.90", "currency": GOOD_CURRENCY_2}, + "type": "fiat", }, ] diff --git a/tests/components/coinbase/test_config_flow.py b/tests/components/coinbase/test_config_flow.py index d153cecc249..fa13648ee71 100644 --- a/tests/components/coinbase/test_config_flow.py +++ b/tests/components/coinbase/test_config_flow.py @@ -19,7 +19,7 @@ from .common import ( mock_get_exchange_rates, mocked_get_accounts, ) -from .const import BAD_CURRENCY, BAD_EXCHANGE_RATE, GOOD_CURRENCY, GOOD_EXCHNAGE_RATE +from .const import BAD_CURRENCY, BAD_EXCHANGE_RATE, GOOD_CURRENCY, GOOD_EXCHANGE_RATE from tests.common import MockConfigEntry @@ -160,7 +160,7 @@ async def test_option_form(hass): result["flow_id"], user_input={ CONF_CURRENCIES: [GOOD_CURRENCY], - CONF_EXCHANGE_RATES: [GOOD_EXCHNAGE_RATE], + CONF_EXCHANGE_RATES: [GOOD_EXCHANGE_RATE], }, ) assert result2["type"] == "create_entry" diff --git a/tests/components/coinbase/test_init.py b/tests/components/coinbase/test_init.py index 36f0ff95472..efb5ba85f73 100644 --- a/tests/components/coinbase/test_init.py +++ b/tests/components/coinbase/test_init.py @@ -3,6 +3,7 @@ from unittest.mock import patch from homeassistant import config_entries from homeassistant.components.coinbase.const import ( + API_TYPE_VAULT, CONF_CURRENCIES, CONF_EXCHANGE_RATES, CONF_YAML_API_TOKEN, @@ -22,8 +23,8 @@ from .common import ( from .const import ( GOOD_CURRENCY, GOOD_CURRENCY_2, - GOOD_EXCHNAGE_RATE, - GOOD_EXCHNAGE_RATE_2, + GOOD_EXCHANGE_RATE, + GOOD_EXCHANGE_RATE_2, ) @@ -34,7 +35,7 @@ async def test_setup(hass): CONF_API_KEY: "123456", CONF_YAML_API_TOKEN: "AbCDeF", CONF_CURRENCIES: [GOOD_CURRENCY, GOOD_CURRENCY_2], - CONF_EXCHANGE_RATES: [GOOD_EXCHNAGE_RATE, GOOD_EXCHNAGE_RATE_2], + CONF_EXCHANGE_RATES: [GOOD_EXCHANGE_RATE, GOOD_EXCHANGE_RATE_2], } } with patch( @@ -54,7 +55,7 @@ async def test_setup(hass): assert entries[0].source == config_entries.SOURCE_IMPORT assert entries[0].options == { CONF_CURRENCIES: [GOOD_CURRENCY, GOOD_CURRENCY_2], - CONF_EXCHANGE_RATES: [GOOD_EXCHNAGE_RATE, GOOD_EXCHNAGE_RATE_2], + CONF_EXCHANGE_RATES: [GOOD_EXCHANGE_RATE, GOOD_EXCHANGE_RATE_2], } @@ -103,7 +104,7 @@ async def test_option_updates(hass: HomeAssistant): result["flow_id"], user_input={ CONF_CURRENCIES: [GOOD_CURRENCY, GOOD_CURRENCY_2], - CONF_EXCHANGE_RATES: [GOOD_EXCHNAGE_RATE, GOOD_EXCHNAGE_RATE_2], + CONF_EXCHANGE_RATES: [GOOD_EXCHANGE_RATE, GOOD_EXCHANGE_RATE_2], }, ) await hass.async_block_till_done() @@ -126,7 +127,7 @@ async def test_option_updates(hass: HomeAssistant): ] assert currencies == [GOOD_CURRENCY, GOOD_CURRENCY_2] - assert rates == [GOOD_EXCHNAGE_RATE, GOOD_EXCHNAGE_RATE_2] + assert rates == [GOOD_EXCHANGE_RATE, GOOD_EXCHANGE_RATE_2] result = await hass.config_entries.options.async_init(config_entry.entry_id) await hass.async_block_till_done() @@ -134,7 +135,7 @@ async def test_option_updates(hass: HomeAssistant): result["flow_id"], user_input={ CONF_CURRENCIES: [GOOD_CURRENCY], - CONF_EXCHANGE_RATES: [GOOD_EXCHNAGE_RATE], + CONF_EXCHANGE_RATES: [GOOD_EXCHANGE_RATE], }, ) await hass.async_block_till_done() @@ -157,4 +158,28 @@ async def test_option_updates(hass: HomeAssistant): ] assert currencies == [GOOD_CURRENCY] - assert rates == [GOOD_EXCHNAGE_RATE] + assert rates == [GOOD_EXCHANGE_RATE] + + +async def test_ignore_vaults_wallets(hass: HomeAssistant): + """Test vaults are ignored in wallet sensors.""" + + with patch( + "coinbase.wallet.client.Client.get_current_user", + return_value=mock_get_current_user(), + ), patch( + "coinbase.wallet.client.Client.get_accounts", new=mocked_get_accounts + ), patch( + "coinbase.wallet.client.Client.get_exchange_rates", + return_value=mock_get_exchange_rates(), + ): + config_entry = await init_mock_coinbase(hass, currencies=[GOOD_CURRENCY]) + await hass.async_block_till_done() + + registry = entity_registry.async_get(hass) + entities = entity_registry.async_entries_for_config_entry( + registry, config_entry.entry_id + ) + assert len(entities) == 1 + entity = entities[0] + assert API_TYPE_VAULT not in entity.original_name.lower()