Skip to content

Commit

Permalink
added unit test for aptos plugin
Browse files Browse the repository at this point in the history
  • Loading branch information
0xaptosj committed Dec 4, 2024
1 parent 608accf commit 1abcd0e
Show file tree
Hide file tree
Showing 3 changed files with 179 additions and 19 deletions.
1 change: 1 addition & 0 deletions packages/plugin-aptos/src/constants.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
export const APT_DECIMALS = 8;
93 changes: 74 additions & 19 deletions packages/plugin-aptos/src/providers/wallet.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,10 @@
import { IAgentRuntime, Memory, Provider, State } from "@ai16z/eliza";
import {
IAgentRuntime,
ICacheManager,
Memory,
Provider,
State,
} from "@ai16z/eliza";
import {
Account,
Aptos,
Expand All @@ -10,6 +16,8 @@ import {
} from "@aptos-labs/ts-sdk";
import BigNumber from "bignumber.js";
import NodeCache from "node-cache";
import * as path from "path";
import { APT_DECIMALS } from "../constants";

// Provider configuration
const PROVIDER_CONFIG = {
Expand All @@ -19,7 +27,7 @@ const PROVIDER_CONFIG = {

interface WalletPortfolio {
totalUsd: string;
totalApt?: string;
totalApt: string;
}

interface Prices {
Expand All @@ -28,15 +36,56 @@ interface Prices {

export class WalletProvider {
private cache: NodeCache;
private cacheKey: string = "aptos/wallet";

constructor(
private aptosClient: Aptos,
private address: string
private address: string,
private cacheManager: ICacheManager
) {
this.cache = new NodeCache({ stdTTL: 300 }); // Cache TTL set to 5 minutes
}

private async fetchAptPriceWithRetry() {
private async readFromCache<T>(key: string): Promise<T | null> {
const cached = await this.cacheManager.get<T>(
path.join(this.cacheKey, key)
);
return cached;
}

private async writeToCache<T>(key: string, data: T): Promise<void> {
await this.cacheManager.set(path.join(this.cacheKey, key), data, {
expires: Date.now() + 5 * 60 * 1000,
});
}

private async getCachedData<T>(key: string): Promise<T | null> {
// Check in-memory cache first
const cachedData = this.cache.get<T>(key);
if (cachedData) {
return cachedData;
}

// Check file-based cache
const fileCachedData = await this.readFromCache<T>(key);
if (fileCachedData) {
// Populate in-memory cache
this.cache.set(key, fileCachedData);
return fileCachedData;
}

return null;
}

private async setCachedData<T>(cacheKey: string, data: T): Promise<void> {
// Set in-memory cache
this.cache.set(cacheKey, data);

// Write to file-based cache
await this.writeToCache(cacheKey, data);
}

private async fetchPricesWithRetry() {
let lastError: Error;

for (let i = 0; i < PROVIDER_CONFIG.MAX_RETRIES; i++) {
Expand Down Expand Up @@ -77,19 +126,20 @@ export class WalletProvider {
async fetchPortfolioValue(): Promise<WalletPortfolio> {
try {
const cacheKey = `portfolio-${this.address}`;
const cachedValue = this.cache.get<WalletPortfolio>(cacheKey);
const cachedValue =
await this.getCachedData<WalletPortfolio>(cacheKey);

if (cachedValue) {
console.log("Cache hit for fetchPortfolioValue");
console.log("Cache hit for fetchPortfolioValue", cachedValue);
return cachedValue;
}
console.log("Cache miss for fetchPortfolioValue");

const aptPrice = await this.fetchAptPrice().catch((error) => {
const prices = await this.fetchPrices().catch((error) => {
console.error("Error fetching APT price:", error);
throw error;
});
const aptAmount = await this.aptosClient
const aptAmountOnChain = await this.aptosClient
.getAccountAPTAmount({
accountAddress: this.address,
})
Expand All @@ -98,32 +148,36 @@ export class WalletProvider {
throw error;
});

const totalUsd = new BigNumber(aptAmount).times(aptPrice.apt.usd);
const aptAmount = new BigNumber(aptAmountOnChain).div(
new BigNumber(10).pow(APT_DECIMALS)
);
const totalUsd = new BigNumber(aptAmount).times(prices.apt.usd);

const portfolio = {
totalUsd: totalUsd.toString(),
totalApt: aptAmount.toString(),
};
this.cache.set(cacheKey, portfolio);
this.setCachedData(cacheKey, portfolio);
console.log("Fetched portfolio:", portfolio);
return portfolio;
} catch (error) {
console.error("Error fetching portfolio:", error);
throw error;
}
}

async fetchAptPrice(): Promise<Prices> {
async fetchPrices(): Promise<Prices> {
try {
const cacheKey = "prices";
const cachedValue = this.cache.get<Prices>(cacheKey);
const cachedValue = await this.getCachedData<Prices>(cacheKey);

if (cachedValue) {
console.log("Cache hit for fetchPrices");
return cachedValue;
}
console.log("Cache miss for fetchPrices");

const aptPriceData = await this.fetchAptPriceWithRetry().catch(
const aptPriceData = await this.fetchPricesWithRetry().catch(
(error) => {
console.error("Error fetching APT price:", error);
throw error;
Expand All @@ -132,7 +186,7 @@ export class WalletProvider {
const prices: Prices = {
apt: { usd: aptPriceData.pair.priceUsd },
};
this.cache.set(cacheKey, prices);
this.setCachedData(cacheKey, prices);
return prices;
} catch (error) {
console.error("Error fetching prices:", error);
Expand All @@ -141,13 +195,13 @@ export class WalletProvider {
}

formatPortfolio(runtime, portfolio: WalletPortfolio): string {
let output = `${runtime.character.description}\n`;
output += `Wallet Address: ${this.address}\n\n`;
let output = `${runtime.character.name}\n`;
output += `Wallet Address: ${this.address}\n`;

const totalUsdFormatted = new BigNumber(portfolio.totalUsd).toFixed(2);
const totalAptFormatted = portfolio.totalApt;
const totalAptFormatted = new BigNumber(portfolio.totalApt).toFixed(4);

output += `Total Value: $${totalUsdFormatted} (${totalAptFormatted} APT)\n\n`;
output += `Total Value: $${totalUsdFormatted} (${totalAptFormatted} APT)\n`;

return output;
}
Expand Down Expand Up @@ -188,7 +242,8 @@ const walletProvider: Provider = {
);
const provider = new WalletProvider(
aptosClient,
aptosAccount.accountAddress.toStringLong()
aptosAccount.accountAddress.toStringLong(),
runtime.cacheManager
);
return await provider.getFormattedPortfolio(runtime);
} catch (error) {
Expand Down
104 changes: 104 additions & 0 deletions packages/plugin-aptos/src/tests/wallet.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
import { describe, it, expect, beforeEach, vi, afterEach } from "vitest";
import { WalletProvider } from "../providers/wallet.ts";
import {
Account,
Aptos,
AptosConfig,
Ed25519PrivateKey,
Network,
PrivateKey,
PrivateKeyVariants,
} from "@aptos-labs/ts-sdk";
import { defaultCharacter } from "@ai16z/eliza";
import BigNumber from "bignumber.js";
import { APT_DECIMALS } from "../constants.ts";

// Mock NodeCache
vi.mock("node-cache", () => {
return {
default: vi.fn().mockImplementation(() => ({
set: vi.fn(),
get: vi.fn().mockReturnValue(null),
})),
};
});

// Mock path module
vi.mock("path", async () => {
const actual = await vi.importActual("path");
return {
...actual,
join: vi.fn().mockImplementation((...args) => args.join("/")),
};
});

// Mock the ICacheManager
const mockCacheManager = {
get: vi.fn().mockResolvedValue(null),
set: vi.fn(),
delete: vi.fn(),
};

describe("WalletProvider", () => {
let walletProvider;
let mockedRuntime;

beforeEach(() => {
vi.clearAllMocks();
mockCacheManager.get.mockResolvedValue(null);

const aptosClient = new Aptos(
new AptosConfig({
network: Network.TESTNET,
})
);
const aptosAccount = Account.fromPrivateKey({
privateKey: new Ed25519PrivateKey(
PrivateKey.formatPrivateKey(
// this is a testnet private key
"0x90e02bf2439492bd9be1ec5f569704accefd65ba88a89c4dcef1977e0203211e",
PrivateKeyVariants.Ed25519
)
),
});

// Create new instance of TokenProvider with mocked dependencies
walletProvider = new WalletProvider(
aptosClient,
aptosAccount.accountAddress.toStringLong(),
mockCacheManager
);

mockedRuntime = {
character: defaultCharacter,
};
});

afterEach(() => {
vi.clearAllTimers();
});

describe("Wallet Integration", () => {
it("should check wallet address", async () => {
const result =
await walletProvider.getFormattedPortfolio(mockedRuntime);

const prices = await walletProvider.fetchPrices();
const aptAmountOnChain =
await walletProvider.aptosClient.getAccountAPTAmount({
accountAddress: walletProvider.address,
});
const aptAmount = new BigNumber(aptAmountOnChain)
.div(new BigNumber(10).pow(APT_DECIMALS))
.toFixed(4);
const totalUsd = new BigNumber(aptAmount)
.times(prices.apt.usd)
.toFixed(2);

expect(result).toEqual(
`Eliza\nWallet Address: ${walletProvider.address}\n` +
`Total Value: $${totalUsd} (${aptAmount} APT)\n`
);
});
});
});

0 comments on commit 1abcd0e

Please sign in to comment.