From 1a66eebe72dc24c5808e249a71746b06da7f8f5c Mon Sep 17 00:00:00 2001
From: Dennis <10233439+idea404@users.noreply.github.com>
Date: Fri, 10 Nov 2023 18:12:45 +0100
Subject: [PATCH] feat: working distribution example

---
 contracts/PensionAccount.sol | 21 +++++++++++++++++++++
 test/main.test.ts            | 30 +++++++++++++++++++++++++++++-
 test/utils.ts                |  6 ++++--
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/contracts/PensionAccount.sol b/contracts/PensionAccount.sol
index 79c1d3e..b151de2 100644
--- a/contracts/PensionAccount.sol
+++ b/contracts/PensionAccount.sol
@@ -25,6 +25,10 @@ contract PensionAccount is IAccount {
     // Expiry timestamp
     uint256 public expiryTimestamp;
 
+    // State variables to track investments
+    uint256 public totalEthReceived;
+    mapping(address => uint256) public investmentPerToken;
+
     // Event for swap action
     event Swap(address indexed token, uint256 amountToSwap);
 
@@ -48,12 +52,29 @@ contract PensionAccount is IAccount {
 
     receive() external payable {
         uint256 amountToSwap = msg.value / 4;
+        totalEthReceived += msg.value; // Track total ETH received
+
+        // Update investment for each token
+        investmentPerToken[DOGE] += amountToSwap;
+        investmentPerToken[PEPE] += amountToSwap;
+        investmentPerToken[SHIB] += amountToSwap;
+        investmentPerToken[BTC] += amountToSwap;
+
+        // Emit Swap events (optional, can be removed if events are not supported)
         emit Swap(DOGE, amountToSwap);
         emit Swap(PEPE, amountToSwap);
         emit Swap(SHIB, amountToSwap);
         emit Swap(BTC, amountToSwap);
     }
 
+    // View function to get investment details
+    function getInvestmentDetails() external view returns (uint256 ethReceived, uint256 dogeInvestment, uint256 pepeInvestment, uint256 shibInvestment, uint256 btcInvestment) {
+        ethReceived = totalEthReceived;
+        dogeInvestment = investmentPerToken[DOGE];
+        pepeInvestment = investmentPerToken[PEPE];
+        shibInvestment = investmentPerToken[SHIB];
+        btcInvestment = investmentPerToken[BTC];
+    }
 
     // Override the executeTransaction function to include the time lock
     function executeTransaction(
diff --git a/test/main.test.ts b/test/main.test.ts
index 5996eff..1c38b35 100644
--- a/test/main.test.ts
+++ b/test/main.test.ts
@@ -113,7 +113,35 @@ describe("Account Abstraction Tests", function () {
       it("Should have a balance", async function () {
         const result = await pensionAccountContract.provider.getBalance(pensionAccountContract.address);
         const balance = parseFloat(ethers.utils.formatEther(result));
-        expect(balance).to.be.greaterThan(99.99);
+        expect(balance).to.be.greaterThan(9.99);
+      });
+
+      it("Should distribute investments correctly", async function () {
+        // Send 10 ETH to the pension account
+        const sendAmount = ethers.utils.parseUnits("20", 18);
+        await firstRichWallet.transfer({
+          to: pensionAccountContract.address,
+          amount: sendAmount,
+          overrides: { type: 113 },
+        });
+  
+        // Call getInvestmentDetails to get investment distribution
+        const investmentDetails = await pensionAccountContract.getInvestmentDetails();
+  
+        // Check total ETH received
+        const ethReceived = parseFloat(ethers.utils.formatEther(investmentDetails.ethReceived));
+        expect(ethReceived).to.equal(120);
+  
+        // Check distribution to each token (30 ETH each)
+        const expectedInvestmentPerToken = 30;
+        const dogeInvestment = parseFloat(ethers.utils.formatEther(investmentDetails.dogeInvestment));
+        expect(dogeInvestment).to.equal(expectedInvestmentPerToken);
+        const pepeInvestment = parseFloat(ethers.utils.formatEther(investmentDetails.pepeInvestment));
+        expect(pepeInvestment).to.equal(expectedInvestmentPerToken);
+        const shibInvestment = parseFloat(ethers.utils.formatEther(investmentDetails.shibInvestment));
+        expect(shibInvestment).to.equal(expectedInvestmentPerToken);
+        const btcInvestment = parseFloat(ethers.utils.formatEther(investmentDetails.btcInvestment));
+        expect(btcInvestment).to.equal(expectedInvestmentPerToken);
       });
     });
   });
diff --git a/test/utils.ts b/test/utils.ts
index 522d43b..e7a2278 100644
--- a/test/utils.ts
+++ b/test/utils.ts
@@ -43,6 +43,8 @@ export async function deployMultisig(wallet: Wallet, factoryAddress: string, own
 
 export async function deployPension(wallet: Wallet, factoryAddress: string, walletOwner: Wallet) {
   const paFactoryArtifact = await hre.artifacts.readArtifact("PensionAccountFactory");
+  const accountArtifact = await hre.artifacts.readArtifact("PensionAccount");
+  
   const paFactory = new ethers.Contract(factoryAddress, paFactoryArtifact.abi, wallet);
 
   // Contract constructor args
@@ -61,14 +63,14 @@ export async function deployPension(wallet: Wallet, factoryAddress: string, wall
 
   // Getting the address of the deployed contract account
   const abiCoder = new ethers.utils.AbiCoder();
-  let multisigAddress = utils.create2Address(
+  let contractAddress = utils.create2Address(
     factoryAddress,
     await paFactory.pensionAccountBytecodeHash(),
     salt,
     abiCoder.encode(["address", "address", "address", "address", "address", "address"], [walletOwner.address, dex, doge, pepe, shib, btc])
   );
 
-  const pensionAccountContract = new ethers.Contract(multisigAddress, paFactoryArtifact.abi, wallet);
+  const pensionAccountContract = new ethers.Contract(contractAddress, accountArtifact.abi, wallet);
   return pensionAccountContract;
 }