Added tests and fixed code for calculating profits and costs

src/RaidGeld.sol

@@ -2,7 +2,8 @@
 pragma solidity ^0.8.13;
 
 import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
-import "@openzeppelin/contracts/access/ownable.sol";
+import "@openzeppelin/contracts/access/Ownable.sol";
+import { RaidGeldUtils } from "../src/RaidGeldUtils.sol";
 
 struct Raider {
     uint16 level;
@@ -13,6 +14,7 @@ struct Army {
     Raider apprentice;
     Raider annointed;
     Raider champion;
+    uint256 profit_per_second;
 }
 
 struct Player {
@@ -52,7 +54,8 @@ contract RaidGeld is ERC20, Ownable {
             moloch_denier: Raider({level: 0}),
             apprentice: Raider({level: 0}),
             annointed: Raider({level: 0}),
-            champion: Raider({level: 0})
+            champion: Raider({level: 0}),
+            profit_per_second: 0
         });
     }
 
@@ -69,8 +72,9 @@ contract RaidGeld is ERC20, Ownable {
     // Manual minting for itchy fingers
     function raid() external onlyPlayer {
         require(block.timestamp >= players[msg.sender].last_raided_at + RAID_WAIT, "Tried minting too soon");
-        // TODO: Make real calculation based on army
+
-        uint256 new_geld = 50 * 10 ** decimals();
+        uint256 time_past = block.timestamp - players[msg.sender].last_raided_at;
+        uint256 new_geld = armies[msg.sender].profit_per_second * time_past * 10 ** decimals();
 
         // TODO: Pink noise, make it so sometimes its better than expected
 
@@ -98,24 +102,43 @@ contract RaidGeld is ERC20, Ownable {
     function addUnit(uint8 unit, uint16 n_units) external onlyPlayer {
         require(unit <= 3, "Unknown unit");
 
-        // TODO: Calculate actual price
+        Army storage army = armies[msg.sender];
-        uint256 price = 1 * 10 ** decimals();
+        uint16 currentLevel = 0;
-        uint256 cost = price * n_units;
+        if (unit == 0) {
-        
+            // moloch_denier
-        require(balanceOf(msg.sender) > cost, "Not enough GELD to add that many units");
+            currentLevel = army.moloch_denier.level;
+        } else if (unit == 1) {
+            // apprentice
+            currentLevel = army.apprentice.level;
+        } else if (unit == 2) {
+            // annointed
+            currentLevel = army.annointed.level;
+        } else if (unit == 3) {
+            // champion
+            currentLevel = army.champion.level;
+        }
+
+        uint256 cost = RaidGeldUtils.calculateUnitPrice(unit, currentLevel, n_units);
+        require(balanceOf(msg.sender) > cost, "Not enough GELD to add this much");
         _burn(msg.sender, cost);
 
-        Army storage army = armies[msg.sender];
         // Increase level
-        if (unit == 0) { // moloch_denier
+        if (unit == 0) {
-           army.moloch_denier.level += n_units;
+            // moloch_denier
-        } else if (unit == 1) { // apprentice
+            army.moloch_denier.level += n_units;
-           army.apprentice.level += n_units;
+        } else if (unit == 1) {
-        } else if (unit == 2) { // annointed
+            // apprentice
-             army.annointed.level += n_units;
+            army.apprentice.level += n_units;
-        } else if (unit == 3) { // champion
+        } else if (unit == 2) {
-           army.champion.level += n_units;
+            // annointed
+            army.annointed.level += n_units;
+        } else if (unit == 3) {
+            // champion
+            army.champion.level += n_units;
         }
+
+        // update profite per second
+        army.profit_per_second = RaidGeldUtils.calculateProfitsPerSecond(army);
     }
 
     receive() external payable {

src/RaidGeldUtils.sol

@@ -1,28 +1,31 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.13;
 
-import {Army} from "../src/RaidGeld.sol" 
+import {Army} from "../src/RaidGeld.sol";
 
 library RaidGeldUtils {
-  function calculateUnitPrice(uint8 unit, uint16 units) internal pure returns (uint256) {
+    function calculateUnitPrice(uint8 unit, uint16 currentLevel, uint16 units) internal pure returns (uint256) {
-    require(unit <= 3, "No matching unit found");
+        require(unit <= 3, "No matching unit found");
-    uint256 price = uint256(unit);
+        uint256 rollingPriceCalculation = uint256(unit) * 38;
-    uint256 unitsAlready = uint256(units);
+        uint256 price = 0;
 
-    // Each level costs 15% more than previous
+        // Each level costs 15% more than previous
-    uint256 PERCENT_INCREASE = 115;
+        uint256 PERCENT_INCREASE = 115;
-    for (uint256 i = 0; i < n; i++) {
+        for (uint256 i = 1; i < currentLevel + units; i++) {
-        newPrice = newPrice * PERCENT_INCREASE / 100;
+            rollingPriceCalculation = rollingPriceCalculation * PERCENT_INCREASE / 100;
+            if (i > currentLevel) {
+                price += rollingPriceCalculation;
+            }
+        }
+        return price;
     }
-    return unitBaseCost + (unitBaseCost * 15 / 100);
-  }
 
-  function calculateProfitsPerSecond(Army army) internal pure returns (uint256) {
+    function calculateProfitsPerSecond(Army memory army) internal pure returns (uint256) {
-    // Each next unit scales progressivelly better
+        // Each next unit scales progressivelly better
-    uint256 moloch_denier_profit = army.moloch_denier.level;
+        uint256 moloch_denier_profit = army.moloch_denier.level;
-    uint256 apprentice_profit = army.apprentice.level * 61 / 10;
+        uint256 apprentice_profit = army.apprentice.level * 61 / 10;
-    uint256 annointed_profit = army.annointed.level * 6 * 64 / 10;
+        uint256 annointed_profit = army.annointed.level * 6 * 64 / 10;
-    uint256 champion_profit = army.champion.level * 61 / 10 * 64 / 10 * 67 / 10;
+        uint256 champion_profit = army.champion.level * 61 / 10 * 64 / 10 * 67 / 10;
-    return moloch_denier_profit + apprentice_profit + annointed_profit + champion_profit;
+        return moloch_denier_profit + apprentice_profit + annointed_profit + champion_profit;
-  }
+    }
 }

test/RaidGeld.t.sol

@@ -28,7 +28,7 @@ contract raid_geldTest is Test {
     function test_00_no_fallback() public {
         vm.expectRevert();
         // Send Ether with some data to trigger fallback
-        address(raid_geld).call{value: 0.1 ether}("0x1234");
+        (bool success, ) = address(raid_geld).call{value: 0.1 ether}("0x1234");
     }
 
     function test_01_no_receive() public {
@@ -65,7 +65,7 @@ contract raid_geldTest is Test {
         assertEq(army.champion.level, 0);
     }
 
-    function test_03_fundsCanBeWithdrawn() public {
+    function test_03_funds_can_be_withdrawn() public {
         uint256 initialBalance = owner.balance;
 
         // Switch to Player 1 and register it
@@ -84,7 +84,7 @@ contract raid_geldTest is Test {
         assertEq(newBalance, initialBalance + raid_geld.BUY_IN_AMOUNT());
     }
 
-    function test_04_onlyOwnerCanWithdraw() public {
+    function test_04_only_owner_can_withdraw() public {
         // Register player 1
         vm.startPrank(player1);
         registerPlayer();
@@ -94,13 +94,58 @@ contract raid_geldTest is Test {
         raid_geld.withdraw();
     }
 
-    function test_05_raid() public {
+    function test_05_is_registered() public {
+        bool is_registered = raid_geld.isRegistered(player1);
+        assertEq(is_registered, false);
+        vm.startPrank(player1);
+        registerPlayer();
+        is_registered = raid_geld.isRegistered(player1);
+        assertEq(is_registered, true);
+    }
+
+    function test_06_add_unit() public {
+        vm.startPrank(player1);
+        registerPlayer();
+
+        vm.expectRevert();
+        // units should be in range
+        raid_geld.addUnit(100, 1);
+
+        // Player should have enough tokens to burn to get units
+        vm.expectRevert();
+        raid_geld.addUnit(1, 100);
+
+        // Add 2 units
+        Army memory army = raid_geld.getArmy(player1);
+        uint256 unit_level = army.annointed.level;
+        uint256 balance = raid_geld.balanceOf(player1);
+        uint256 income_per_sec = army.profit_per_second;  
+        raid_geld.addUnit(0, 2);
+
+        // TODO: maybe try to test against exact values we want here
+
+        // Check that those tokens were burnt
+        uint256 newBalance = raid_geld.balanceOf(player1);
+        assertGt(newBalance, balance - 2 * 10 ** 4);
+        army = raid_geld.getArmy(player1);
+
+        // Check that unit level increased
+        uint256 new_unit_level = army.moloch_denier.level;
+        assertEq(new_unit_level, unit_level + 2);
+
+        // Check that user income per second increased
+        uint256 new_income_per_sec = army.profit_per_second;  
+        assertLt(income_per_sec, new_income_per_sec);
+    }
+
+    function test_07_raid() public {
         // Let some time pass so we dont start at block timestamp 0
         vm.warp(120);
 
         // Register player 1
         vm.startPrank(player1);
         registerPlayer();
+        raid_geld.addUnit(1, 2);
 
         uint256 balance = raid_geld.balanceOf(player1);
 
@@ -128,42 +173,4 @@ contract raid_geldTest is Test {
         assertLt(newBalance, newestBalance);
         assertLt(last_raided_at, last_raided_at_2);
     }
-
-    function test_06_is_registered() public {
-        bool is_registered = raid_geld.isRegistered(player1);
-        assertEq(is_registered, false);
-        vm.startPrank(player1);
-        registerPlayer();
-        is_registered = raid_geld.isRegistered(player1);
-        assertEq(is_registered, true);
-    }
-
-    function test_07_add_unit() public {
-        vm.startPrank(player1);
-        registerPlayer();
-
-        vm.expectRevert();
-        // units should be in range
-        raid_geld.addUnit(100, 1);
-
-        // Player should have enough tokens to burn to get units
-        vm.expectRevert();
-        raid_geld.addUnit(1, 100);
-
-        // Add 2 units
-        Army memory army = raid_geld.getArmy(player1);
-        uint256 unit_level = army.annointed.level;
-        uint256 balance = raid_geld.balanceOf(player1);
-        raid_geld.addUnit(2, 2);
-        
-        // Check that those tokens were burnt
-        // TODO: Correct price
-        uint256 newBalance = raid_geld.balanceOf(player1);
-        assertEq(newBalance, balance - 2 * 10 ** 4);
-        army = raid_geld.getArmy(player1);
-
-        // Check that unit level increased
-        uint256 new_unit_level = army.annointed.level;
-        assertEq(new_unit_level, unit_level + 2);
-    }
 }