From 44c89b4b6be4ccde475ee6afba325277619e0f35 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Leszczyk?= <michal.leszczyk@avalabs.org>
Date: Thu, 20 Feb 2025 09:39:22 +0100
Subject: [PATCH] fix: balance polling when closed

---
 .../balances/BalancePollingService.test.ts    | 26 +++++++++++++++++++
 .../balances/BalancePollingService.ts         | 12 +++++----
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/background/services/balances/BalancePollingService.test.ts b/src/background/services/balances/BalancePollingService.test.ts
index 52c70327c..ffcfd0de8 100644
--- a/src/background/services/balances/BalancePollingService.test.ts
+++ b/src/background/services/balances/BalancePollingService.test.ts
@@ -36,6 +36,32 @@ describe('src/background/services/balances/BalancePollingService.ts', () => {
     }
   };
 
+  it('stops polling when extension is closed', async () => {
+    const service = new BalancePollingService(aggregatorServiceMock);
+
+    await service.startPolling(
+      account,
+      activeNetworkId,
+      roundRobinChainIds,
+      tokenTypes,
+    );
+
+    // Purposefully DO NOT await this, as we want to test the behavior
+    // of onAllExtensionsClosed() being called while polling is in progress.
+    runIntervalTimes(1);
+
+    expect(aggregatorServiceMock.getBalancesForNetworks).toHaveBeenCalledTimes(
+      2, // Once immediately after startPolling() and one more time after the first interval
+    );
+
+    service.onAllExtensionsClosed();
+
+    await runIntervalTimes(10); // Wait X intervals, 10 is arbitrary here
+    expect(aggregatorServiceMock.getBalancesForNetworks).toHaveBeenCalledTimes(
+      2, // Still only called twice. No more polling after onAllExtensionsClosed()
+    );
+  });
+
   describe('when polling is active', () => {
     beforeEach(async () => {
       const service = new BalancePollingService(aggregatorServiceMock);
diff --git a/src/background/services/balances/BalancePollingService.ts b/src/background/services/balances/BalancePollingService.ts
index a6f71c57d..7cd79cc43 100644
--- a/src/background/services/balances/BalancePollingService.ts
+++ b/src/background/services/balances/BalancePollingService.ts
@@ -15,10 +15,7 @@ export class BalancePollingService implements OnLock, OnAllExtensionClosed {
   #timer: NodeJS.Timeout | null = null;
   #pollingIteration = 0;
   #lastPollingStartedAt?: number;
-
-  get isPollingActive() {
-    return this.#timer !== null;
-  }
+  #isPollingActive = false;
 
   constructor(private balanceAggregator: BalanceAggregatorService) {}
 
@@ -35,6 +32,7 @@ export class BalancePollingService implements OnLock, OnAllExtensionClosed {
     // Stop any polling that may be occurring already
     this.stopPolling();
     // Start a new interval
+    this.#isPollingActive = true;
     return this.pollBalances(
       account,
       activeChainId,
@@ -48,6 +46,7 @@ export class BalancePollingService implements OnLock, OnAllExtensionClosed {
   }
 
   stopPolling() {
+    this.#isPollingActive = false;
     if (this.#timer) {
       clearTimeout(this.#timer);
       this.#timer = null;
@@ -85,7 +84,10 @@ export class BalancePollingService implements OnLock, OnAllExtensionClosed {
 
     // Only schedule the next update if another polling was not started
     // while we were waiting for balance results.
-    if (thisPollingStartedAt === this.#lastPollingStartedAt) {
+    if (
+      this.#isPollingActive &&
+      thisPollingStartedAt === this.#lastPollingStartedAt
+    ) {
       this.scheduleNextUpdate(
         account,
         activeChainId,