Skip to content

Commit 46e5f65

Browse files
gboudreaugboudreau
authored
Merge pull request #34 from ebastos/main
User dictionary lookup for easy maintenance of account types
2 parents 68a2b69 + d9c6497 commit 46e5f65

3 files changed

Lines changed: 106 additions & 48 deletions

File tree

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
44

55
[project]
66
name = "ws-api"
7-
version = "0.29.1"
7+
version = "0.29.2"
88
description = "Access your Wealthsimple account using their (GraphQL) API."
99
readme = "README.md"
1010
license = {file = "LICENSE"}

‎tests/test_wealthsimple_api_happy.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,72 @@ def test_account_add_description_cash(api):
174174
assert account["number"] == "cust123"
175175

176176

177+
def test_account_add_description_cash_joint(api):
178+
"""Test CASH type with MULTI_OWNER returns 'Cash: joint'."""
179+
account = {
180+
"id": "acc1",
181+
"unifiedAccountType": "CASH",
182+
"accountOwnerConfiguration": "MULTI_OWNER",
183+
"accountFeatures": [],
184+
"custodianAccounts": [],
185+
}
186+
api._account_add_description(account)
187+
assert account["description"] == "Cash: joint"
188+
189+
190+
def test_account_add_description_managed_private_credit(api):
191+
"""Test MANAGED_NON_REGISTERED with PRIVATE_CREDIT feature."""
192+
account = {
193+
"id": "acc1",
194+
"unifiedAccountType": "MANAGED_NON_REGISTERED",
195+
"accountOwnerConfiguration": "SINGLE_OWNER",
196+
"accountFeatures": [{"name": "PRIVATE_CREDIT"}],
197+
"custodianAccounts": [],
198+
}
199+
api._account_add_description(account)
200+
assert account["description"] == "Non-registered: managed - private credit"
201+
202+
203+
def test_account_add_description_managed_private_equity(api):
204+
"""Test MANAGED_NON_REGISTERED with PRIVATE_EQUITY feature."""
205+
account = {
206+
"id": "acc1",
207+
"unifiedAccountType": "MANAGED_NON_REGISTERED",
208+
"accountOwnerConfiguration": "SINGLE_OWNER",
209+
"accountFeatures": [{"name": "PRIVATE_EQUITY"}],
210+
"custodianAccounts": [],
211+
}
212+
api._account_add_description(account)
213+
assert account["description"] == "Non-registered: managed - private equity"
214+
215+
216+
def test_account_add_description_nickname_override(api):
217+
"""Test that nickname takes precedence over unifiedAccountType."""
218+
account = {
219+
"id": "acc1",
220+
"unifiedAccountType": "SELF_DIRECTED_RRSP",
221+
"nickname": "My Retirement Fund",
222+
"accountOwnerConfiguration": "SINGLE_OWNER",
223+
"accountFeatures": [],
224+
"custodianAccounts": [],
225+
}
226+
api._account_add_description(account)
227+
assert account["description"] == "My Retirement Fund"
228+
229+
230+
def test_account_add_description_unknown_type_fallback(api):
231+
"""Test that unknown account types fall back to unifiedAccountType value."""
232+
account = {
233+
"id": "acc1",
234+
"unifiedAccountType": "SOME_FUTURE_TYPE",
235+
"accountOwnerConfiguration": "SINGLE_OWNER",
236+
"accountFeatures": [],
237+
"custodianAccounts": [],
238+
}
239+
api._account_add_description(account)
240+
assert account["description"] == "SOME_FUTURE_TYPE"
241+
242+
177243
def test_get_account_balances(api):
178244
"""Test get_account_balances happy path."""
179245
fake_accounts = [

‎ws_api/wealthsimple_api.py‎

Lines changed: 39 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,26 @@
1717
)
1818
from ws_api.session import WSAPISession
1919

20+
# Mapping of account types to human-readable descriptions
21+
_ACCOUNT_TYPE_DESCRIPTIONS = {
22+
"SELF_DIRECTED_RRSP": "RRSP: self-directed",
23+
"MANAGED_RRSP": "RRSP: managed",
24+
"SELF_DIRECTED_SPOUSAL_RRSP": "RRSP: self-directed spousal",
25+
"SELF_DIRECTED_TFSA": "TFSA: self-directed",
26+
"MANAGED_TFSA": "TFSA: managed",
27+
"SELF_DIRECTED_FHSA": "FHSA: self-directed",
28+
"MANAGED_FHSA": "FHSA: managed",
29+
"SELF_DIRECTED_NON_REGISTERED": "Non-registered: self-directed",
30+
"SELF_DIRECTED_JOINT_NON_REGISTERED": "Non-registered: self-directed - joint",
31+
"SELF_DIRECTED_NON_REGISTERED_MARGIN": "Non-registered: self-directed margin",
32+
"MANAGED_JOINT": "Non-registered: managed - joint",
33+
"SELF_DIRECTED_CRYPTO": "Crypto",
34+
"SELF_DIRECTED_RRIF": "RRIF: self-directed",
35+
"SELF_DIRECTED_SPOUSAL_RRIF": "RRIF: self-directed spousal",
36+
"CREDIT_CARD": "Credit card",
37+
"SELF_DIRECTED_LIRA": "LIRA: self-directed",
38+
}
39+
2040

2141
class WealthsimpleAPIBase:
2242
OAUTH_BASE_URL = "https://api.production.wealthsimple.com/v1/oauth/v2"
@@ -442,64 +462,36 @@ def _account_add_description(account):
442462
account["number"] = account["id"]
443463
# This is the account number visible in the WS app:
444464
for ca in account["custodianAccounts"]:
445-
if (ca["branch"] in ["WS", "TR"]) and ca["status"] == "open":
465+
if ca["branch"] in ("WS", "TR") and ca["status"] == "open":
446466
account["number"] = ca["id"]
447467

448-
# Default
449-
account["description"] = account["unifiedAccountType"]
450-
451468
if account.get("nickname"):
452469
account["description"] = account["nickname"]
453-
elif account["unifiedAccountType"] == "CASH":
470+
return
471+
472+
account_type = account["unifiedAccountType"]
473+
474+
# Special case: CASH depends on owner configuration
475+
if account_type == "CASH":
454476
account["description"] = (
455477
"Cash: joint"
456478
if account["accountOwnerConfiguration"] == "MULTI_OWNER"
457479
else "Cash"
458480
)
459-
elif account["unifiedAccountType"] == "SELF_DIRECTED_RRSP":
460-
account["description"] = "RRSP: self-directed"
461-
elif account["unifiedAccountType"] == "MANAGED_RRSP":
462-
account["description"] = "RRSP: managed"
463-
elif account["unifiedAccountType"] == "SELF_DIRECTED_SPOUSAL_RRSP":
464-
account["description"] = "RRSP: self-directed spousal"
465-
elif account["unifiedAccountType"] == "SELF_DIRECTED_TFSA":
466-
account["description"] = "TFSA: self-directed"
467-
elif account["unifiedAccountType"] == "MANAGED_TFSA":
468-
account["description"] = "TFSA: managed"
469-
elif account["unifiedAccountType"] == "SELF_DIRECTED_FHSA":
470-
account["description"] = f"FHSA: self-directed"
471-
elif account["unifiedAccountType"] == "MANAGED_FHSA":
472-
account["description"] = f"FHSA: managed"
473-
elif account["unifiedAccountType"] == "SELF_DIRECTED_NON_REGISTERED":
474-
account["description"] = "Non-registered: self-directed"
475-
elif account["unifiedAccountType"] == "SELF_DIRECTED_JOINT_NON_REGISTERED":
476-
account["description"] = "Non-registered: self-directed - joint"
477-
elif account["unifiedAccountType"] == "SELF_DIRECTED_NON_REGISTERED_MARGIN":
478-
account["description"] = "Non-registered: self-directed margin"
479-
elif account["unifiedAccountType"] == "MANAGED_JOINT":
480-
account["description"] = "Non-registered: managed - joint"
481-
elif account["unifiedAccountType"] == "SELF_DIRECTED_CRYPTO":
482-
account["description"] = "Crypto"
483-
elif account["unifiedAccountType"] == "SELF_DIRECTED_RRIF":
484-
account["description"] = "RRIF: self-directed"
485-
elif account["unifiedAccountType"] == "SELF_DIRECTED_SPOUSAL_RRIF":
486-
account["description"] = "RRIF: self-directed spousal"
487-
elif account["unifiedAccountType"] == "CREDIT_CARD":
488-
account["description"] = "Credit card"
489-
elif account["unifiedAccountType"] == "MANAGED_NON_REGISTERED":
490-
if any(
491-
feature["name"] == "PRIVATE_CREDIT"
492-
for feature in account["accountFeatures"]
493-
):
481+
# Special case: MANAGED_NON_REGISTERED depends on features
482+
elif account_type == "MANAGED_NON_REGISTERED":
483+
features = {f["name"] for f in account["accountFeatures"]}
484+
if "PRIVATE_CREDIT" in features:
494485
account["description"] = "Non-registered: managed - private credit"
495-
if any(
496-
feature["name"] == "PRIVATE_EQUITY"
497-
for feature in account["accountFeatures"]
498-
):
486+
elif "PRIVATE_EQUITY" in features:
499487
account["description"] = "Non-registered: managed - private equity"
500-
elif account["unifiedAccountType"] == "SELF_DIRECTED_LIRA":
501-
account["description"] = "LIRA: self-directed"
502-
# TODO: Add other types as needed
488+
else:
489+
account["description"] = account_type
490+
# Simple lookup for all other types
491+
else:
492+
account["description"] = _ACCOUNT_TYPE_DESCRIPTIONS.get(
493+
account_type, account_type
494+
)
503495

504496
def get_account_balances(self, account_id):
505497
accounts = self.do_graphql_query(

0 commit comments

Comments
 (0)