From 19b9454c924b1a683799dd21495a5d7ddbded142 Mon Sep 17 00:00:00 2001
From: Dmitry Tatarkin <tatarkin@gmail.com>
Date: Wed, 25 Sep 2024 12:28:18 +0300
Subject: [PATCH] Fix `Account` pydantic model to properly handle serialization
 and deserialization

---
 alpaca/broker/models/accounts.py   | 50 ------------------------------
 tests/broker/factories/accounts.py | 20 ++++++++++++
 tests/broker/test_models.py        | 10 ++++++
 3 files changed, 30 insertions(+), 50 deletions(-)

diff --git a/alpaca/broker/models/accounts.py b/alpaca/broker/models/accounts.py
index a067afd2..f1275c7f 100644
--- a/alpaca/broker/models/accounts.py
+++ b/alpaca/broker/models/accounts.py
@@ -262,56 +262,6 @@ class Account(ModelWithID):
     documents: Optional[List[AccountDocument]] = None
     trusted_contact: Optional[TrustedContact] = None
 
-    def __init__(self, **response):
-        super().__init__(
-            id=(UUID(response["id"])),
-            account_number=(response["account_number"]),
-            status=(response["status"]),
-            crypto_status=(
-                response["crypto_status"] if "crypto_status" in response else None
-            ),
-            kyc_results=(
-                TypeAdapter(KycResults).validate_python(response["kyc_results"])
-                if "kyc_results" in response and response["kyc_results"] is not None
-                else None
-            ),
-            currency=(response["currency"]),
-            last_equity=(response["last_equity"]),
-            created_at=(response["created_at"]),
-            contact=(
-                TypeAdapter(Contact).validate_python(response["contact"])
-                if "contact" in response
-                else None
-            ),
-            identity=(
-                TypeAdapter(Identity).validate_python(response["identity"])
-                if "identity" in response
-                else None
-            ),
-            disclosures=(
-                TypeAdapter(Disclosures).validate_python(response["disclosures"])
-                if "disclosures" in response
-                else None
-            ),
-            agreements=(
-                TypeAdapter(List[Agreement]).validate_python(response["agreements"])
-                if "agreements" in response
-                else None
-            ),
-            documents=(
-                TypeAdapter(List[AccountDocument]).validate_python(
-                    response["documents"]
-                )
-                if "documents" in response
-                else None
-            ),
-            trusted_contact=(
-                TypeAdapter(TrustedContact).validate_python(response["trusted_contact"])
-                if "trusted_contact" in response
-                else None
-            ),
-        )
-
 
 class TradeAccount(BaseTradeAccount):
     """
diff --git a/tests/broker/factories/accounts.py b/tests/broker/factories/accounts.py
index 9bad062d..4971473b 100644
--- a/tests/broker/factories/accounts.py
+++ b/tests/broker/factories/accounts.py
@@ -1,4 +1,6 @@
 from typing import List
+
+from alpaca.broker import Account
 from alpaca.trading.enums import DTBPCheck, PDTCheck
 from alpaca.trading.models import AccountConfiguration as TradeAccountConfiguration
 
@@ -143,3 +145,21 @@ def create_dummy_trade_account_configuration() -> TradeAccountConfiguration:
         trade_confirm_email="all",
         ptp_no_exception_entry=False,
     )
+
+
+def create_dummy_account() -> Account:
+    """
+    Create a basic account instance with prefilled data
+
+    Returns:
+        Account: a prefilled Account instance for testing
+    """
+
+    return Account(
+        id="2d6cab28-c5d1-4ff8-91c6-b6404a9ee114",
+        account_number="551081356",
+        status="ACTIVE",
+        currency="USD",
+        last_equity="0",
+        created_at="2024-07-10T18:35:52.244506Z",
+    )
diff --git a/tests/broker/test_models.py b/tests/broker/test_models.py
index 311aa229..8864bdf1 100644
--- a/tests/broker/test_models.py
+++ b/tests/broker/test_models.py
@@ -1,5 +1,8 @@
 from datetime import datetime
 
+from alpaca.broker import Account
+from tests.broker.factories import accounts as factory
+
 import pytest
 
 from alpaca.broker.requests import (
@@ -466,3 +469,10 @@ def test_journal_with_amount_and_qty():
         )
 
     assert "Cash journals must contain an amount to transfer." in str(e.value)
+
+
+def test_account_serialization_deserialization():
+    account = factory.create_dummy_account()
+    json_str = account.model_dump_json()
+    account2 = Account.model_validate_json(json_str)
+    assert account2 == account