prepare for multiple withdrawal outputs in asset balance calculation

MixinNetwork · Nov 6, 2024 · 74f09d4 · 74f09d4
1 parent ae8c252
commit 74f09d4
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 1 deletion.
diff --git a/storage/badger_asset.go b/storage/badger_asset.go
@@ -54,7 +54,11 @@ func writeTotalInAsset(txn *badger.Txn, ver *common.VersionedTransaction) error
 	typ := ver.TransactionType()
 	switch { // TODO needs full test code for all kind of transactions
 	case typ == common.TransactionTypeWithdrawalSubmit:
-		total = total.Sub(ver.Outputs[0].Amount)
+		for _, o := range ver.Outputs {
+			if o.Type == common.OutputTypeWithdrawalSubmit {
+				total = total.Sub(o.Amount)
+			}
+		}
 	case typ == common.TransactionTypeDeposit:
 		total = total.Add(ver.DepositData().Amount)
 	case typ == common.TransactionTypeMint:

diff --git a/storage/transaction_test.go b/storage/transaction_test.go
@@ -28,6 +28,7 @@ func TestTransaction(t *testing.T) {
 	require.Equal(uint64(0), seq)
 
 	gns, err := common.ReadGenesis("../config/genesis.json")
+	require.Nil(err)
 	rounds, snapshots, transactions, err := gns.BuildSnapshots()
 	require.Nil(err)
 	loaded, err := store.CheckGenesisLoad(snapshots)
@@ -40,6 +41,12 @@ func TestTransaction(t *testing.T) {
 	require.True(loaded)
 	signers := []crypto.Hash{rounds[0].NodeId}
 
+	genesis := common.NewInteger(13439).Mul(27).Add(common.NewInteger(2700))
+	_, balance, err := store.ReadAssetWithBalance(common.XINAssetId)
+	require.Nil(err)
+	require.Equal(genesis.String(), balance.String())
+	require.Equal("365553.00000000", balance.String())
+
 	seed := make([]byte, 64)
 	crypto.ReadRand(seed)
 	mixin := common.NewAddressFromSeed(seed)
@@ -72,6 +79,9 @@ func TestTransaction(t *testing.T) {
 	utxo, err := store.ReadUTXOLock(deposit.AsVersioned().PayloadHash(), 0)
 	require.Nil(err)
 	require.Nil(utxo)
+	_, balance, err = store.ReadAssetWithBalance(common.XINAssetId)
+	require.Nil(err)
+	require.Equal("365553.00000000", balance.String())
 
 	round, _ := store.ReadRound(rounds[0].NodeId)
 	require.Equal(uint64(1), round.Number)
@@ -92,6 +102,9 @@ func TestTransaction(t *testing.T) {
 	utxo, err = store.ReadUTXOLock(deposit.AsVersioned().PayloadHash(), 0)
 	require.Nil(err)
 	require.NotNil(utxo)
+	_, balance, err = store.ReadAssetWithBalance(common.XINAssetId)
+	require.Nil(err)
+	require.Equal("365563.00000000", balance.String())
 
 	submit := common.NewTransactionV5(common.XINAssetId)
 	submit.AddInput(deposit.AsVersioned().PayloadHash(), 0)
@@ -107,6 +120,9 @@ func TestTransaction(t *testing.T) {
 	require.Nil(err)
 	err = store.WriteTransaction(submit.AsVersioned())
 	require.Nil(err)
+	_, balance, err = store.ReadAssetWithBalance(common.XINAssetId)
+	require.Nil(err)
+	require.Equal("365563.00000000", balance.String())
 
 	snap = &common.Snapshot{
 		Version:      common.SnapshotVersionCommonEncoding,
@@ -122,6 +138,9 @@ func TestTransaction(t *testing.T) {
 	}
 	err = store.WriteSnapshot(topo, signers)
 	require.Nil(err)
+	_, balance, err = store.ReadAssetWithBalance(common.XINAssetId)
+	require.Nil(err)
+	require.Equal("365562.00000000", balance.String())
 
 	ver, ss, err := store.ReadWithdrawalClaim(submit.AsVersioned().PayloadHash())
 	require.Nil(err)
@@ -159,4 +178,8 @@ func TestTransaction(t *testing.T) {
 	require.Nil(err)
 	require.Equal(topo.PayloadHash().String(), ss)
 	require.Equal(claim.AsVersioned().PayloadHash(), ver.PayloadHash())
+
+	_, balance, err = store.ReadAssetWithBalance(common.XINAssetId)
+	require.Nil(err)
+	require.Equal("365562.00000000", balance.String())
 }