fix: handle compute unit + lookup tables, and improve flaky tests (#965)

* fix: handle compute unit + lookup tables, and improve flaky tests

* more complex: multiple instructions + multiple indices

pkg/solana/client/client_test.go

@@ -278,14 +278,12 @@ func TestClient_GetBlocks(t *testing.T) {
 
 	// Verify we can retrieve blocks
 	startSlot := uint64(1)
-	endSlot := uint64(6)
+	endSlot := uint64(10)
 	require.Eventually(t,
 		func() bool {
 			blocks, err := c.GetBlocks(ctx, startSlot, &endSlot)
-			if err != nil {
-				return false
-			}
-			return len(blocks) == 5
+			require.NoError(t, err) // don't mask error within false
+			return len(blocks) >= 2 // slots != blocks (expect multiple blocks for 10 slots)
 		},
 		requestTimeout, 500*time.Millisecond)
 }

pkg/solana/fees/computebudget.go

@@ -151,6 +151,17 @@ func set(tx *solana.Transaction, baseData instruction, appendToFront bool) error
 
 		// https://github.com/gagliardetto/solana-go/blob/618f56666078f8131a384ab27afd918d248c08b7/transaction.go#L293
 		tx.Message.Header.NumReadonlyUnsignedAccounts++
+
+		// lookup table addresses are indexed after the tx.Message.AccountKeys
+		// higher indices must be increased
+		// https://github.com/gagliardetto/solana-go/blob/da2193071f56059aa35010a239cece016c4e827f/transaction.go#L440
+		for i, ix := range tx.Message.Instructions {
+			for j, v := range ix.Accounts {
+				if int(v) >= programIdx {
+					tx.Message.Instructions[i].Accounts[j]++
+				}
+			}
+		}
 	}
 
 	// get instruction data

pkg/solana/fees/computebudget_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/gagliardetto/solana-go"
 	"github.com/gagliardetto/solana-go/programs/system"
+	"github.com/gagliardetto/solana-go/programs/token"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -136,6 +137,55 @@ func testSet[V instruction](t *testing.T, builder func(uint) V, setter func(*sol
 		assert.NoError(t, err)
 		assert.Equal(t, data, []byte(tx.Message.Instructions[computeIndex].Data))
 	})
+
+	t.Run("with_lookuptables", func(t *testing.T) {
+		t.Parallel()
+
+		// build base tx (no fee)
+		tx, err := solana.NewTransaction([]solana.Instruction{
+			system.NewTransferInstruction(
+				1,
+				solana.PublicKey{1},
+				solana.PublicKey{2},
+			).Build(),
+			token.NewTransferInstruction(
+				uint64(1),
+				solana.PublicKey{11},
+				solana.PublicKey{12},
+				solana.PublicKey{13},
+				[]solana.PublicKey{
+					solana.PublicKey{14},
+				},
+			).Build(),
+		},
+			solana.Hash{},
+			solana.TransactionAddressTables(map[solana.PublicKey]solana.PublicKeySlice{
+				solana.PublicKey{}: solana.PublicKeySlice{solana.PublicKey{1}, solana.PublicKey{2}, solana.PublicKey{11}, solana.PublicKey{12}, solana.PublicKey{13}, solana.PublicKey{14}},
+			}),
+		)
+		require.NoError(t, err)
+
+		// check current account indices
+		assert.Equal(t, 2, len(tx.Message.Instructions))
+		assert.Equal(t, []uint16{0, 4}, tx.Message.Instructions[0].Accounts)
+		assert.Equal(t, []uint16{5, 6, 7, 1}, tx.Message.Instructions[1].Accounts)
+		assert.Equal(t, 4, len(tx.Message.AccountKeys))
+
+		// add fee
+		require.NoError(t, setter(tx, builder(0)))
+
+		// evaluate
+		assert.Equal(t, 3, len(tx.Message.Instructions))
+		computeUnitIndex := getIndex(len(tx.Message.Instructions))
+		transferIndex := 0
+		if computeUnitIndex == transferIndex {
+			transferIndex = 1
+		}
+		assert.Equal(t, 5, len(tx.Message.AccountKeys))
+		assert.Equal(t, uint16(4), tx.Message.Instructions[computeUnitIndex].ProgramIDIndex)
+		assert.Equal(t, []uint16{0, 5}, tx.Message.Instructions[transferIndex].Accounts)
+		assert.Equal(t, []uint16{6, 7, 8, 1}, tx.Message.Instructions[transferIndex+1].Accounts)
+	})
 }
 
 func TestParse(t *testing.T) {