pump-parser/pumpamm_test.go

package pump_parser

import (
	"bytes"
	"encoding/base64"
	"fmt"

	agbinary "github.com/gagliardetto/binary"
	"github.com/gagliardetto/solana-go"
	"github.com/mr-tron/base58"

	"testing"
)

type Pool struct {
	Bump                  uint8            // 1
	Index                 uint16           // 3
	Creator               solana.PublicKey // 35
	BaseMint              solana.PublicKey // 67
	QuoteMint             solana.PublicKey // 99
	LpMint                solana.PublicKey // 131
	PoolBaseTokenAccount  solana.PublicKey // 163
	PoolQuoteTokenAccount solana.PublicKey // 195
	LpSupply              uint64           // 203
	CoinCreator           solana.PublicKey // 235
	IsMayhemMode          bool             // 236
}

func TestDecodePoolData(t *testing.T) {
	var base64Data = "8ZptBBGxbbz9AAB/+s6++xNrbunrNkEpE9IY0qfortUqEOB5s3kKGSvL+9H0YLudD5fPy6QWG5FDfdl3FQRE3H804Ivg5J2OnSHvBpuIV/6rgYT7aH9jRhjANdrEOdwa6ztVmKDwAAAAAAGcWzgf1bP0MRLCJXxkgXyazv6CfdxoJGAh01hfNVqrkGoU6aOndSrbraRyP929gjNuQ4qg/nj2RjtLlBazO47we2q7UVSLFd69kGDUes9pm4n4dftp/lIZS068RLXhIR1niGtZ0AMAAFC4Q0TJ2zOhAXgWLRoYc4eSeu9nlTcupLKiXLyIDytOAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=="
	data, _ := base64.StdEncoding.DecodeString(base64Data)
	var pool Pool
	fmt.Println("data length:", len(data))
	err := agbinary.NewBorshDecoder(data[8:]).Decode(&pool)
	if err != nil {
		t.Errorf("Failed to decode pool data: %v", err)
	}
	fmt.Println(pool)
}

func TestAmmBuyEvent(t *testing.T) {
	fmt.Println(pumpAmmBuyEventDiscriminator)
	fmt.Println(pumpGetFeesDiscriminator)
}

func TestDecodeAmmBuyEventSuffixCompatibility(t *testing.T) {
	prefix := ammBuyEventPrefix{TimeStamp: 123, PoolQuoteTokenReserve: 456}

	legacy, err := decodeAmmBuyEvent(encodeAmmEventParts(t, prefix))
	if err != nil {
		t.Fatalf("decodeAmmBuyEvent(legacy) error = %v", err)
	}
	if legacy.TimeStamp != prefix.TimeStamp || legacy.PoolQuoteTokenReserve != prefix.PoolQuoteTokenReserve {
		t.Fatalf("legacy prefix mismatch: %+v", legacy.ammBuyEventPrefix)
	}
	if legacy.BuybackFeeBasisPoints != 0 || legacy.BuybackFee != 0 || legacy.CanBoost || legacy.BaseSupply != 0 || !pumpAmmVirtualQuoteReserve(legacy.VirtualQuoteReserves).IsZero() {
		t.Fatalf("legacy suffix is not zero-valued: %+v", legacy)
	}

	buyback := ammTradeEventBuybackSuffix{BuybackFeeBasisPoints: 1000, BuybackFee: 77}
	withBuyback, err := decodeAmmBuyEvent(encodeAmmEventParts(t, prefix, buyback))
	if err != nil {
		t.Fatalf("decodeAmmBuyEvent(buyback) error = %v", err)
	}
	if withBuyback.BuybackFeeBasisPoints != buyback.BuybackFeeBasisPoints || withBuyback.BuybackFee != buyback.BuybackFee {
		t.Fatalf("buyback suffix mismatch: %+v", withBuyback)
	}
	if withBuyback.CanBoost || withBuyback.BaseSupply != 0 || !pumpAmmVirtualQuoteReserve(withBuyback.VirtualQuoteReserves).IsZero() {
		t.Fatalf("pre-virtual suffix has unexpected virtual fields: %+v", withBuyback)
	}

	negativeVirtual := agbinary.Int128(agbinary.Uint128{Lo: ^uint64(122), Hi: ^uint64(0)})
	virtual := ammTradeEventVirtualSuffix{
		VirtualQuoteReserves: negativeVirtual,
		CanBoost:             true,
		BaseSupply:           1_000_000,
	}
	current, err := decodeAmmBuyEvent(encodeAmmEventParts(t, prefix, buyback, virtual))
	if err != nil {
		t.Fatalf("decodeAmmBuyEvent(current) error = %v", err)
	}
	if got := pumpAmmVirtualQuoteReserve(current.VirtualQuoteReserves).String(); got != "-123" {
		t.Fatalf("VirtualQuoteReserves = %s, want -123", got)
	}
	if !current.CanBoost || current.BaseSupply != virtual.BaseSupply {
		t.Fatalf("current suffix mismatch: %+v", current)
	}
}

func TestDecodeAmmSellEventSuffixCompatibility(t *testing.T) {
	prefix := ammSellEventPrefix{Timestamp: 321, PoolQuoteTokenReserves: 654, QuoteAmountOut: 54}
	buyback := ammTradeEventBuybackSuffix{BuybackFeeBasisPoints: 900, BuybackFee: 88}
	virtual := ammTradeEventVirtualSuffix{
		VirtualQuoteReserves: agbinary.Int128(agbinary.Uint128{Lo: 987}),
		CanBoost:             true,
		BaseSupply:           2_000_000,
	}

	event, err := decodeAmmSellEvent(encodeAmmEventParts(t, prefix, buyback, virtual))
	if err != nil {
		t.Fatalf("decodeAmmSellEvent() error = %v", err)
	}
	if event.Timestamp != prefix.Timestamp || event.BuybackFee != buyback.BuybackFee || !event.CanBoost || event.BaseSupply != virtual.BaseSupply {
		t.Fatalf("decoded sell event mismatch: %+v", event)
	}
	if got := pumpAmmVirtualQuoteReserve(event.VirtualQuoteReserves).String(); got != "987" {
		t.Fatalf("VirtualQuoteReserves = %s, want 987", got)
	}
	real, effective := pumpAmmSellPostQuoteReserves(event)
	if real.String() != "600" || effective.String() != "1587" {
		t.Fatalf("post reserves = real %s effective %s", real, effective)
	}
}

func TestDecodeAmmTradeEventRejectsTruncatedSuffix(t *testing.T) {
	prefix := encodeAmmEventParts(t, ammBuyEventPrefix{})
	for _, suffixLength := range []int{1, 15, 17, 40} {
		data := append(append([]byte(nil), prefix...), make([]byte, suffixLength)...)
		if _, err := decodeAmmBuyEvent(data); err == nil {
			t.Fatalf("decodeAmmBuyEvent() error = nil for %d-byte suffix", suffixLength)
		}
	}
}

func TestDecodeAmmBuyEventDevnetVirtualReserve(t *testing.T) {
	const eventDataBase58 = "6MF1ykxMQW5eFmo1ErWgds8V2wfEyqJPjEsoz4yRoKQCktPze5U297neF7huAQQ9DhxEJqHPwezxUh9mcQEsfdmnwW84YVxcfEzDUduRqc1zty2dLQZcFGzFeFC6fmpX75tkvuMvTN89HekHa9TEikGRuXhtWkeacjNoVqfbxXUp3FQxGSFoyeBaSdAfPMvXJ7zN4ALTDtn3a5g7HxJCPPWeUtYGcJig6RpKrC3DzAQ3x5SiiJiVhK1zdocmkkRXuTezt5L5Svrf9pbM4QkkwgvxrYqfUs5YeaLNqWYtNZfimLtENhQJcfrhQh2svCosCJ2Ux5iyyLZwJu7VLZibD41CwxuF6ngmrhde8BmEegmQ9f3L4qYhqgW5RJSiu1JnKMFqmTckuwGydQfTiyA9RsCtFxiazaAVim9SMARU2fqe9zfAPdRjBetAvTGkBxrERmmzhWVhER6g448mV1E1eeQjFWYRhtR33hSUpNUUP6dUF8CWae2BuJxaKexpiY3pyqUwbH7osvexbff8jpfpZGubP5UWjqhwijtF8tdkoEdAnb2TgWXqXMu4by5VDsEzwbgNK1DbkGbxrcXYwFTNskqoUqoNLPRaAcfnjTMA9cVme62HAj5WvxuzpL9Gpw7ZzSE2LYnVvzTa6PnN2bbpuqJpZ3d"
	data, err := base58.Decode(eventDataBase58)
	if err != nil {
		t.Fatalf("base58.Decode() error = %v", err)
	}
	if len(data) < 16 || !bytes.Equal(data[:8], pumpAmmEventDiscriminator[:]) || !bytes.Equal(data[8:16], pumpAmmBuyEventDiscriminator[:]) {
		t.Fatalf("unexpected event discriminator: len=%d", len(data))
	}

	event, err := decodeAmmBuyEvent(data[16:])
	if err != nil {
		t.Fatalf("decodeAmmBuyEvent(devnet) error = %v", err)
	}
	if event.TimeStamp != 1784271481 || event.PoolQuoteTokenReserve != 2_213_159_021 || event.QuoteAmountIn != 100_000_000 || event.QuoteAmountInWithLpFee != 98_785_184 {
		t.Fatalf("devnet buy event core fields mismatch: %+v", event)
	}
	if event.BuybackFeeBasisPoints != 1000 || event.BuybackFee != 91_851 {
		t.Fatalf("devnet buyback fields mismatch: bps=%d fee=%d", event.BuybackFeeBasisPoints, event.BuybackFee)
	}
	if got := pumpAmmVirtualQuoteReserve(event.VirtualQuoteReserves).String(); got != "583150126" {
		t.Fatalf("VirtualQuoteReserves = %s, want 583150126", got)
	}
	if !event.CanBoost || event.BaseSupply != 1_000_000_000_000_000 {
		t.Fatalf("devnet boost metadata mismatch: can_boost=%t base_supply=%d", event.CanBoost, event.BaseSupply)
	}

	real, effective := pumpAmmBuyPostQuoteReserves(event)
	if real.String() != "2313159021" || effective.String() != "2896309147" {
		t.Fatalf("post reserves = real %s effective %s", real, effective)
	}
}

func encodeAmmEventParts(t *testing.T, parts ...any) []byte {
	t.Helper()

	var buf bytes.Buffer
	encoder := agbinary.NewBorshEncoder(&buf)
	for _, part := range parts {
		if err := encoder.Encode(part); err != nil {
			t.Fatalf("borsh encode %T: %v", part, err)
		}
	}
	return buf.Bytes()
}