From 5ccc843b990cb9902c8c07fd8f16936b0d71f19c Mon Sep 17 00:00:00 2001
From: Rafael Cardenas <rafael@rafaelcr.com>
Date: Thu, 3 Oct 2024 09:58:16 -0600
Subject: [PATCH] fix: token amount coalesce

---
 src/datastore/pg-store.ts |  6 +++---
 tests/api/mempool.test.ts | 37 +++++++++++++++++++++++++++++++++----
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/src/datastore/pg-store.ts b/src/datastore/pg-store.ts
index 04619ea11..c1175b7cc 100644
--- a/src/datastore/pg-store.ts
+++ b/src/datastore/pg-store.ts
@@ -2483,17 +2483,17 @@ export class PgStore extends BasePgStore {
   async getPrincipalMempoolStxBalanceDelta(sql: PgSqlClient, principal: string): Promise<bigint> {
     const results = await sql<{ delta: string }[]>`
       WITH sent AS (
-        SELECT token_transfer_amount + COALESCE(fee_rate, 0) AS total
+        SELECT SUM(COALESCE(token_transfer_amount, 0) + fee_rate) AS total
         FROM mempool_txs
         WHERE pruned = false AND sender_address = ${principal}
       ),
       sponsored AS (
-        SELECT COALESCE(fee_rate, 0) AS total
+        SELECT SUM(fee_rate) AS total
         FROM mempool_txs
         WHERE pruned = false AND sponsor_address = ${principal} AND sponsored = true
       ),
       received AS (
-        SELECT token_transfer_amount AS total
+        SELECT SUM(COALESCE(token_transfer_amount, 0)) AS total
         FROM mempool_txs
         WHERE pruned = false AND token_transfer_recipient_address = ${principal}
       )
diff --git a/tests/api/mempool.test.ts b/tests/api/mempool.test.ts
index 8fe496b70..6d2d5ab9f 100644
--- a/tests/api/mempool.test.ts
+++ b/tests/api/mempool.test.ts
@@ -2167,6 +2167,25 @@ describe('mempool tests', () => {
     expect(balance1.body.balance).toEqual('2000');
     expect(balance1.body.estimated_balance).toEqual('1850'); // Minus amount and fee
 
+    // Contract call in mempool
+    await db.updateMempoolTxs({
+      mempoolTxs: [
+        testMempoolTx({
+          tx_id: '0x0002aa',
+          sender_address: address,
+          type_id: DbTxTypeId.ContractCall,
+          token_transfer_amount: 0n,
+          contract_call_contract_id: 'ST27W5M8BRKA7C5MZE2R1S1F4XTPHFWFRNHA9M04Y.hello-world',
+          contract_call_function_args: '',
+          contract_call_function_name: 'test',
+          fee_rate: 50n,
+        }),
+      ],
+    });
+    const balance1b = await supertest(api.server).get(url);
+    expect(balance1b.body.balance).toEqual('2000');
+    expect(balance1b.body.estimated_balance).toEqual('1800'); // Minus fee
+
     // Sponsored tx in mempool
     await db.updateMempoolTxs({
       mempoolTxs: [
@@ -2181,7 +2200,7 @@ describe('mempool tests', () => {
     });
     const balance2 = await supertest(api.server).get(url);
     expect(balance2.body.balance).toEqual('2000');
-    expect(balance2.body.estimated_balance).toEqual('1800'); // Minus fee
+    expect(balance2.body.estimated_balance).toEqual('1750'); // Minus fee
 
     // STX received in mempool
     await db.updateMempoolTxs({
@@ -2196,7 +2215,7 @@ describe('mempool tests', () => {
     });
     const balance3 = await supertest(api.server).get(url);
     expect(balance3.body.balance).toEqual('2000');
-    expect(balance3.body.estimated_balance).toEqual('1900'); // Plus amount
+    expect(balance3.body.estimated_balance).toEqual('1850'); // Plus amount
 
     // Confirm all txs
     await db.update(
@@ -2212,6 +2231,16 @@ describe('mempool tests', () => {
           fee_rate: 50n,
         })
         .addTxStxEvent({ sender: address, amount: 100n })
+        .addTx({
+          tx_id: '0x0002aa',
+          sender_address: address,
+          type_id: DbTxTypeId.ContractCall,
+          token_transfer_amount: 0n,
+          contract_call_contract_id: 'ST27W5M8BRKA7C5MZE2R1S1F4XTPHFWFRNHA9M04Y.hello-world',
+          contract_call_function_args: '',
+          contract_call_function_name: 'test',
+          fee_rate: 50n,
+        })
         .addTx({
           tx_id: '0x0003',
           sponsor_address: address,
@@ -2229,7 +2258,7 @@ describe('mempool tests', () => {
         .build()
     );
     const balance4 = await supertest(api.server).get(url);
-    expect(balance4.body.balance).toEqual('1900');
-    expect(balance4.body.estimated_balance).toEqual('1900');
+    expect(balance4.body.balance).toEqual('1850');
+    expect(balance4.body.estimated_balance).toEqual('1850');
   });
 });