added unit test for aptos plugin

0xaptosj · 0xaptosj · commit 3efd50d44199 · 2024-12-04T02:49:49.000-08:00
diff --git a/packages/plugin-aptos/src/constants.ts b/packages/plugin-aptos/src/constants.ts
@@ -0,0 +1 @@
+export const APT_DECIMALS = 8;
diff --git a/packages/plugin-aptos/src/providers/wallet.ts b/packages/plugin-aptos/src/providers/wallet.ts
@@ -1,4 +1,10 @@
-import { IAgentRuntime, Memory, Provider, State } from "@ai16z/eliza";
+import {
+    IAgentRuntime,
+    ICacheManager,
+    Memory,
+    Provider,
+    State,
+} from "@ai16z/eliza";
 import {
     Account,
     Aptos,
@@ -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 = {
@@ -19,7 +27,7 @@ const PROVIDER_CONFIG = {
 
 interface WalletPortfolio {
     totalUsd: string;
-    totalApt?: string;
+    totalApt: string;
 }
 
 interface Prices {
@@ -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++) {
@@ -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,
                 })
@@ -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;
@@ -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);
@@ -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;
     }
@@ -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) {
diff --git a/packages/plugin-aptos/src/tests/wallet.test.ts b/packages/plugin-aptos/src/tests/wallet.test.ts
@@ -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`
+            );
+        });
+    });
+});
