Skip to content

Commit

Permalink
Add tests for swap
Browse files Browse the repository at this point in the history
  • Loading branch information
@pauldragonfly
pauldragonfly committed Jan 15, 2025
1 parent 33763fb commit 5c30240
Show file tree
Hide file tree
Showing 2 changed files with 342 additions and 11 deletions.
238 changes: 235 additions & 3 deletions ts-sdk/whirlpool/tests/swap.test.ts
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,237 @@
import { describe } from "vitest";
import { fetchToken } from "@solana-program/token-2022";
import type { Address } from "@solana/web3.js";
import assert from "assert";
import { beforeAll, describe, it } from "vitest";
import { swapInstructions } from "../src/swap";
import { rpc, sendTransaction } from "./utils/mockRpc";
import {
setupPosition,
setupTEPosition,
setupWhirlpool,
} from "./utils/program";
import { setupAta, setupMint } from "./utils/token";

describe.skip("Swap", () => {
// TODO: <-
import {
setupAtaTE,
setupMintTE,
setupMintTEFee,
} from "./utils/tokenExtensions";

const mintTypes = new Map([
["A", setupMint],
["B", setupMint],
["TEA", setupMintTE],
["TEB", setupMintTE],
["TEFee", setupMintTEFee],
]);

const ataTypes = new Map([
["A", setupAta],
["B", setupAta],
["TEA", setupAtaTE],
["TEB", setupAtaTE],
["TEFee", setupAtaTE],
]);

const poolTypes = new Map([
["A-B", setupWhirlpool],
["A-TEA", setupWhirlpool],
["TEA-TEB", setupWhirlpool],
["A-TEFee", setupWhirlpool],
]);

const positionTypes = new Map([
["equally centered", { tickLower: -100, tickUpper: 100 }],
["one sided A", { tickLower: -100, tickUpper: -1 }],
["one sided B", { tickLower: 1, tickUpper: 100 }],
]);

describe("Swap", () => {
const atas: Map<string, Address> = new Map();
const initialLiquidity = 100_000n;
const mints: Map<string, Address> = new Map();
const pools: Map<string, Address> = new Map();
const positions: Map<string, Address> = new Map();
const tickSpacing = 64;
const tokenBalance = 1_000_000n;

beforeAll(async () => {
for (const [name, setup] of mintTypes) {
mints.set(name, await setup());
}

for (const [name, setup] of ataTypes) {
const mint = mints.get(name)!;
atas.set(name, await setup(mint, { amount: tokenBalance }));
}

for (const [name, setup] of poolTypes) {
const [mintAKey, mintBKey] = name.split("-");
const mintA = mints.get(mintAKey)!;
const mintB = mints.get(mintBKey)!;
pools.set(name, await setup(mintA, mintB, tickSpacing));
}

for (const [poolName, poolAddress] of pools) {
for (const [positionTypeName, tickRange] of positionTypes) {
const position = await setupPosition(poolAddress, {
...tickRange,
liquidity: initialLiquidity,
});
positions.set(`${poolName} ${positionTypeName}`, position);

const positionTE = await setupTEPosition(poolAddress, {
...tickRange,
liquidity: initialLiquidity,
});
positions.set(`TE ${poolName} ${positionTypeName}`, positionTE);
}
}
});

const testSwapAExactIn = async (poolName: string) => {
const [mintAName, mintBName] = poolName.split("-");
const mintAAddress = mints.get(mintAName)!;
const ataAAddress = atas.get(mintAName)!;
const ataBAddress = atas.get(mintBName)!;
const poolAddress = pools.get(poolName)!;

let tokenABefore = await fetchToken(rpc, ataAAddress);
let tokenBBefore = await fetchToken(rpc, ataBAddress);

const { instructions, quote } = await swapInstructions(
rpc,
{ inputAmount: 100n, mint: mintAAddress },
poolAddress,
100, // slippage
);
await sendTransaction(instructions);

let tokenAAfter = await fetchToken(rpc, ataAAddress);
let tokenBAfter = await fetchToken(rpc, ataBAddress);

assert.strictEqual(
-quote.tokenIn,
tokenAAfter.data.amount - tokenABefore.data.amount,
);

assert.strictEqual(
quote.tokenEstOut,
tokenBAfter.data.amount - tokenBBefore.data.amount,
);
};

const testSwapAExactOut = async (poolName: string) => {
const [mintAName, mintBName] = poolName.split("-");
const mintAAddress = mints.get(mintAName)!;
const ataAAddress = atas.get(mintAName)!;
const ataBAddress = atas.get(mintBName)!;
const poolAddress = pools.get(poolName)!;

let tokenABefore = await fetchToken(rpc, ataAAddress);
let tokenBBefore = await fetchToken(rpc, ataBAddress);

const { instructions, quote } = await swapInstructions(
rpc,
{ outputAmount: 100n, mint: mintAAddress },
poolAddress,
100, // slippage
);
await sendTransaction(instructions);

let tokenAAfter = await fetchToken(rpc, ataAAddress);
let tokenBAfter = await fetchToken(rpc, ataBAddress);

assert.strictEqual(
quote.tokenOut,
tokenAAfter.data.amount - tokenABefore.data.amount,
);

assert.strictEqual(
-quote.tokenEstIn,
tokenBAfter.data.amount - tokenBBefore.data.amount,
);
};

const testSwapBExactIn = async (poolName: string) => {
const [mintAName, mintBName] = poolName.split("-");
const mintBAddress = mints.get(mintBName)!;
const ataAAddress = atas.get(mintAName)!;
const ataBAddress = atas.get(mintBName)!;
const poolAddress = pools.get(poolName)!;

let tokenABefore = await fetchToken(rpc, ataAAddress);
let tokenBBefore = await fetchToken(rpc, ataBAddress);

const { instructions, quote } = await swapInstructions(
rpc,
{ inputAmount: 100n, mint: mintBAddress },
poolAddress,
100, // slippage
);
await sendTransaction(instructions);

let tokenAAfter = await fetchToken(rpc, ataAAddress);
let tokenBAfter = await fetchToken(rpc, ataBAddress);

assert.strictEqual(
quote.tokenEstOut,
tokenAAfter.data.amount - tokenABefore.data.amount,
);

assert.strictEqual(
-quote.tokenIn,
tokenBAfter.data.amount - tokenBBefore.data.amount,
);
};

const testSwapBExactOut = async (poolName: string) => {
const [mintAName, mintBName] = poolName.split("-");
const mintBAddress = mints.get(mintBName)!;
const ataAAddress = atas.get(mintAName)!;
const ataBAddress = atas.get(mintBName)!;
const poolAddress = pools.get(poolName)!;

let tokenABefore = await fetchToken(rpc, ataAAddress);
let tokenBBefore = await fetchToken(rpc, ataBAddress);

const { instructions, quote } = await swapInstructions(
rpc,
{ outputAmount: 100n, mint: mintBAddress },
poolAddress,
100, // slippage
);
await sendTransaction(instructions);

let tokenAAfter = await fetchToken(rpc, ataAAddress);
let tokenBAfter = await fetchToken(rpc, ataBAddress);

assert.strictEqual(
-quote.tokenEstIn,
tokenAAfter.data.amount - tokenABefore.data.amount,
);

assert.strictEqual(
quote.tokenOut,
tokenBAfter.data.amount - tokenBBefore.data.amount,
);
};

for (const poolName of poolTypes.keys()) {
it(`Should swap A to B in ${poolName} using A amount`, async () => {
await testSwapAExactIn(poolName);
});

it(`Should swap B to A in ${poolName} using A amount`, async () => {
await testSwapAExactOut(poolName);
});

it(`Should swap B to A in ${poolName} using B amount`, async () => {
await testSwapBExactIn(poolName);
});

it(`Should swap A to B in ${poolName} using B amount`, async () => {
await testSwapBExactOut(poolName);
});
}
});
115 changes: 107 additions & 8 deletions ts-sdk/whirlpool/tests/utils/program.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import {
fetchAllMaybeTickArray,
fetchWhirlpool,
getFeeTierAddress,
getIncreaseLiquidityV2Instruction,
getInitializeConfigInstruction,
getInitializeFeeTierInstruction,
getInitializePoolV2Instruction,
Expand All @@ -13,25 +14,27 @@ import {
getTokenBadgeAddress,
getWhirlpoolAddress,
} from "@orca-so/whirlpools-client";
import { address, type Address, type IInstruction } from "@solana/web3.js";
import { rpc, sendTransaction, signer } from "./mockRpc";
import {
SPLASH_POOL_TICK_SPACING,
WHIRLPOOLS_CONFIG_ADDRESS,
} from "../../src/config";
import {
getInitializableTickIndex,
getTickArrayStartTickIndex,
increaseLiquidityQuote,
tickIndexToSqrtPrice,
} from "@orca-so/whirlpools-core";
import { MEMO_PROGRAM_ADDRESS } from "@solana-program/memo";
import { TOKEN_PROGRAM_ADDRESS } from "@solana-program/token";
import {
TOKEN_2022_PROGRAM_ADDRESS,
ASSOCIATED_TOKEN_PROGRAM_ADDRESS,
fetchMint,
findAssociatedTokenPda,
TOKEN_2022_PROGRAM_ADDRESS,
} from "@solana-program/token-2022";
import { address, type Address, type IInstruction } from "@solana/web3.js";
import {
SPLASH_POOL_TICK_SPACING,
WHIRLPOOLS_CONFIG_ADDRESS,
} from "../../src/config";
import { getNextKeypair } from "./keypair";
import { TOKEN_PROGRAM_ADDRESS } from "@solana-program/token";
import { rpc, sendTransaction, signer } from "./mockRpc";

export async function setupConfigAndFeeTiers(): Promise<Address> {
const keypair = getNextKeypair();
Expand Down Expand Up @@ -232,6 +235,54 @@ export async function setupPosition(
}),
);

if (config.liquidity) {
const tokenMintA = await fetchMint(rpc, whirlpoolAccount.data.tokenMintA);
const tokenOwnerAccountA = await findAssociatedTokenPda({
owner: signer.address,
mint: whirlpoolAccount.data.tokenMintA,
tokenProgram: tokenMintA.programAddress,
}).then((x) => x[0]);

const tokenMintB = await fetchMint(rpc, whirlpoolAccount.data.tokenMintB);
const tokenOwnerAccountB = await findAssociatedTokenPda({
owner: signer.address,
mint: whirlpoolAccount.data.tokenMintB,
tokenProgram: tokenMintB.programAddress,
}).then((x) => x[0]);

const quote = increaseLiquidityQuote(
config.liquidity,
100,
whirlpoolAccount.data.sqrtPrice,
initializableLowerTickIndex,
initializableUpperTickIndex,
);

instructions.push(
getIncreaseLiquidityV2Instruction({
whirlpool: whirlpool,
positionAuthority: signer,
position: positionAddress[0],
positionTokenAccount,
tokenOwnerAccountA: tokenOwnerAccountA,
tokenOwnerAccountB: tokenOwnerAccountB,
tokenVaultA: whirlpoolAccount.data.tokenVaultA,
tokenVaultB: whirlpoolAccount.data.tokenVaultB,
tokenMintA: whirlpoolAccount.data.tokenMintA,
tokenMintB: whirlpoolAccount.data.tokenMintB,
tokenProgramA: tokenMintA.programAddress,
tokenProgramB: tokenMintB.programAddress,
tickArrayLower: lowerTickArrayAddress,
tickArrayUpper: upperTickArrayAddress,
liquidityAmount: quote.liquidityDelta,
tokenMaxA: quote.tokenMaxA,
tokenMaxB: quote.tokenMaxB,
memoProgram: MEMO_PROGRAM_ADDRESS,
remainingAccountsInfo: null,
}),
);
}

await sendTransaction(instructions);

return positionMint.address;
Expand Down Expand Up @@ -331,6 +382,54 @@ export async function setupTEPosition(
}),
);

if (config.liquidity) {
const tokenMintA = await fetchMint(rpc, whirlpoolAccount.data.tokenMintA);
const tokenOwnerAccountA = await findAssociatedTokenPda({
owner: signer.address,
mint: whirlpoolAccount.data.tokenMintA,
tokenProgram: tokenMintA.programAddress,
}).then((x) => x[0]);

const tokenMintB = await fetchMint(rpc, whirlpoolAccount.data.tokenMintB);
const tokenOwnerAccountB = await findAssociatedTokenPda({
owner: signer.address,
mint: whirlpoolAccount.data.tokenMintB,
tokenProgram: tokenMintB.programAddress,
}).then((x) => x[0]);

const quote = increaseLiquidityQuote(
config.liquidity,
100,
whirlpoolAccount.data.sqrtPrice,
initializableLowerTickIndex,
initializableUpperTickIndex,
);

instructions.push(
getIncreaseLiquidityV2Instruction({
whirlpool: whirlpool,
positionAuthority: signer,
position: positionAddress[0],
positionTokenAccount,
tokenOwnerAccountA: tokenOwnerAccountA,
tokenOwnerAccountB: tokenOwnerAccountB,
tokenVaultA: whirlpoolAccount.data.tokenVaultA,
tokenVaultB: whirlpoolAccount.data.tokenVaultB,
tokenMintA: whirlpoolAccount.data.tokenMintA,
tokenMintB: whirlpoolAccount.data.tokenMintB,
tokenProgramA: tokenMintA.programAddress,
tokenProgramB: tokenMintB.programAddress,
tickArrayLower: lowerTickArrayAddress,
tickArrayUpper: upperTickArrayAddress,
liquidityAmount: quote.liquidityDelta,
tokenMaxA: quote.tokenMaxA,
tokenMaxB: quote.tokenMaxB,
memoProgram: MEMO_PROGRAM_ADDRESS,
remainingAccountsInfo: null,
}),
);
}

await sendTransaction(instructions);

return positionMint.address;
Expand Down

0 comments on commit 5c30240

Please sign in to comment.