Merge pull request #30 from tuxcanfly/fix-update-balance-genesis

Fix update balance genesis

internal/storage/block_storage.go

@@ -490,14 +490,22 @@ func (b *BlockStorage) UpdateBalance(
 	}
 
 	var existingValue string
-	if exists {
+	switch {
+	case exists:
+		// This could happen if balances are bootstrapped and should not be
+		// overridden.
 		parseBal, err := parseBalanceEntry(balance)
 		if err != nil {
 			return err
 		}
 
 		existingValue = parseBal.Amount.Value
-	} else {
+	case parentBlock != nil && change.Block.Hash == parentBlock.Hash:
+		// Don't attempt to use the helper if we are going to query the same
+		// block we are processing (causes the duplicate issue).
+		existingValue = "0"
+	default:
+		// Use helper to fetch existing balance.
 		amount, err := b.helper.AccountBalance(ctx, change.Account, change.Currency, parentBlock)
 		if err != nil {
 			return fmt.Errorf("%w: unable to get previous account balance", err)

internal/storage/block_storage_test.go

@@ -251,6 +251,9 @@ func TestBlock(t *testing.T) {
 
 func TestBalance(t *testing.T) {
 	var (
+		genesisAccount = &types.AccountIdentifier{
+			Address: "genesis",
+		}
 		account = &types.AccountIdentifier{
 			Address: "blah",
 		}
@@ -323,6 +326,10 @@ func TestBalance(t *testing.T) {
 		amountNilCurrency = &types.Amount{
 			Value: "100",
 		}
+		genesisBlock = &types.BlockIdentifier{
+			Hash:  "0",
+			Index: 0,
+		}
 		newBlock = &types.BlockIdentifier{
 			Hash:  "kdasdj",
 			Index: 123890,
@@ -343,7 +350,11 @@ func TestBalance(t *testing.T) {
 			Value:    "-1000",
 			Currency: currency,
 		}
-		mockHelper = &MockBlockStorageHelper{}
+		mockHelper = &MockBlockStorageHelper{
+			AccountBalances: map[string]string{
+				"genesis": "100",
+			},
+		}
 	)
 
 	ctx := context.Background()
@@ -368,6 +379,31 @@ func TestBalance(t *testing.T) {
 		assert.Equal(t, newBlock, block)
 	})
 
+	t.Run("Set and get genesis balance", func(t *testing.T) {
+		txn := storage.newDatabaseTransaction(ctx, true)
+		err := storage.UpdateBalance(
+			ctx,
+			txn,
+			&parser.BalanceChange{
+				Account:    genesisAccount,
+				Currency:   currency,
+				Block:      genesisBlock,
+				Difference: amount.Value,
+			},
+			genesisBlock,
+		)
+		assert.NoError(t, err)
+		assert.NoError(t, txn.Commit(ctx))
+
+		amount, block, err := storage.GetBalance(ctx, genesisAccount, currency, genesisBlock)
+		assert.NoError(t, err)
+		assert.Equal(t, &types.Amount{
+			Value:    "100",
+			Currency: currency,
+		}, amount)
+		assert.Equal(t, genesisBlock, block)
+	})
+
 	t.Run("Set and get balance", func(t *testing.T) {
 		txn := storage.newDatabaseTransaction(ctx, true)
 		err := storage.UpdateBalance(
@@ -603,6 +639,10 @@ func TestBalance(t *testing.T) {
 		accounts, err := storage.GetAllAccountCurrency(ctx)
 		assert.NoError(t, err)
 		assert.ElementsMatch(t, []*reconciler.AccountCurrency{
+			{
+				Account:  genesisAccount,
+				Currency: currency,
+			},
 			{
 				Account:  account,
 				Currency: currency,
@@ -823,6 +863,7 @@ func TestCreateBlockCache(t *testing.T) {
 
 type MockBlockStorageHelper struct {
 	AccountBalanceAmount string
+	AccountBalances      map[string]string
 	ExemptAccounts       []*reconciler.AccountCurrency
 }
 
@@ -832,6 +873,13 @@ func (h *MockBlockStorageHelper) AccountBalance(
 	currency *types.Currency,
 	block *types.BlockIdentifier,
 ) (*types.Amount, error) {
+	if balance, ok := h.AccountBalances[account.Address]; ok {
+		return &types.Amount{
+			Value:    balance,
+			Currency: currency,
+		}, nil
+	}
+
 	value := "0"
 	if len(h.AccountBalanceAmount) > 0 {
 		value = h.AccountBalanceAmount