From cc38e4555a799f549e2e5d428f2e967210db064f Mon Sep 17 00:00:00 2001
From: Alex Gherghisan <alexg@aztecprotocol.com>
Date: Mon, 6 Nov 2023 12:47:30 +0000
Subject: [PATCH] fix: register account waits for it to be synched

---
 .../aztec.js/src/account/manager/index.ts     | 30 ++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/yarn-project/aztec.js/src/account/manager/index.ts b/yarn-project/aztec.js/src/account/manager/index.ts
index 69f73ebea8d..0e3a53c65c9 100644
--- a/yarn-project/aztec.js/src/account/manager/index.ts
+++ b/yarn-project/aztec.js/src/account/manager/index.ts
@@ -8,6 +8,7 @@ import {
   DeployMethod,
   WaitOpts,
   generatePublicKey,
+  retryUntil,
 } from '../../index.js';
 import { AccountContract, Salt } from '../index.js';
 import { AccountInterface } from '../interface.js';
@@ -88,11 +89,13 @@ export class AccountManager {
    * Registers this account in the PXE Service and returns the associated wallet. Registering
    * the account on the PXE Service is required for managing private state associated with it.
    * Use the returned wallet to create Contract instances to be interacted with from this account.
+   * @param opts - Options to wait for the account to be registered.
    * @returns A Wallet instance.
    */
-  public async register(): Promise<AccountWalletWithPrivateKey> {
+  public async register(opts: WaitOpts = {}): Promise<AccountWalletWithPrivateKey> {
     const completeAddress = await this.getCompleteAddress();
     await this.pxe.registerAccount(this.encryptionPrivateKey, completeAddress.partialAddress);
+    await this.waitSynch(opts);
     return this.getWallet();
   }
 
@@ -142,4 +145,29 @@ export class AccountManager {
     await this.deploy().then(tx => tx.wait(opts));
     return this.getWallet();
   }
+
+  /**
+   * Waits for the account to finish synchronizing with the PXE Service.
+   * @param opts - Options to wait for account to finish synchronizing
+   * @returns A wallet instance
+   */
+  public async waitSynch({ interval, timeout }: WaitOpts): Promise<AccountWalletWithPrivateKey> {
+    const address = (await this.getCompleteAddress()).address;
+    await retryUntil(
+      async () => {
+        const status = await this.pxe.getSyncStatus();
+        const accountSynchedToBlock = status.notes[address.toString()];
+        if (!accountSynchedToBlock) {
+          return false;
+        } else {
+          return accountSynchedToBlock >= status.blocks;
+        }
+      },
+      'waitSynch',
+      interval,
+      timeout,
+    );
+
+    return this.getWallet();
+  }
 }