Added army upgrading, removed total amount as that part is tracked by openzeppelin ERC20

src/RaidGeld.sol

@@ -4,9 +4,6 @@ pragma solidity ^0.8.13;
 import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import "@openzeppelin/contracts/access/Ownable.sol";
 
-// DEV
-import "forge-std/console.sol";
-
 struct Raider {
     uint16 level;
 }
@@ -29,7 +26,6 @@ contract RaidGeld is ERC20, Ownable {
     uint256 public constant BUY_IN_AMOUNT = 0.0005 ether;
     uint256 public constant INITIAL_GELD = 50 * 10 ** DECIMALS;
     uint256 public constant RAID_WAIT = 15 seconds;
-    uint256 public total_minted = 0;
 
     mapping(address => Player) private players;
     mapping(address => Army) private armies;
@@ -58,7 +54,6 @@ contract RaidGeld is ERC20, Ownable {
             annointed: Raider({level: 0}),
             champion: Raider({level: 0})
         });
-        total_minted += INITIAL_GELD;
     }
 
     // Override for default number of decimals
@@ -76,11 +71,12 @@ contract RaidGeld is ERC20, Ownable {
         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();
+
+        // TODO: Pink noise, make it so sometimes its better than expected
+
         _mint(msg.sender, new_geld);
-        total_minted += new_geld;
-        console.log(players[msg.sender].last_raided_at);
         players[msg.sender].last_raided_at = block.timestamp;
-        console.log(players[msg.sender].last_raided_at);
+        players[msg.sender].total_minted = new_geld;
     }
 
     // Function to get Player struct
@@ -98,6 +94,30 @@ contract RaidGeld is ERC20, Ownable {
         return players[addr].created_at != 0;
     }
 
+    // Add a unit to your army
+    function addUnit(uint8 unit, uint16 n_units) external onlyPlayer {
+        require(unit <= 3, "Unknown unit");
+
+        // TODO: Calculate actual price
+        uint256 price = 1 * 10 ** decimals();
+        uint256 cost = price * n_units;
+        
+        require(balanceOf(msg.sender) > cost, "Not enough GELD to add that many units");
+        _burn(msg.sender, cost);
+
+        Army storage army = armies[msg.sender];
+        // Increase level
+        if (unit == 0) { // moloch_denier
+           army.moloch_denier.level += n_units;
+        } else if (unit == 1) { // apprentice
+           army.apprentice.level += n_units;
+        } else if (unit == 2) { // annointed
+             army.annointed.level += n_units;
+        } else if (unit == 3) { // champion
+           army.champion.level += n_units;
+        }
+    }
+
     receive() external payable {
         revert("No plain Ether accepted, use register() function to check in :)");
     }

test/RaidGeld.t.sol

@@ -41,7 +41,6 @@ contract raid_geldTest is Test {
         vm.startPrank(player1);
 
         uint256 initialBalance = address(raid_geld).balance;
-        uint256 initialTotalMinted = raid_geld.total_minted();
 
         // Send registration fee ETH to the contract
         registerPlayer();
@@ -64,9 +63,6 @@ contract raid_geldTest is Test {
         assertEq(army.apprentice.level, 0);
         assertEq(army.annointed.level, 0);
         assertEq(army.champion.level, 0);
-
-        // Verify that total_minted is updated
-        assertEq(raid_geld.total_minted(), initialTotalMinted + raid_geld.INITIAL_GELD());
     }
 
     function test_03_fundsCanBeWithdrawn() public {
@@ -106,21 +102,16 @@ contract raid_geldTest is Test {
         vm.startPrank(player1);
         registerPlayer();
 
-
         uint256 balance = raid_geld.balanceOf(player1);
-        uint256 total_minted = raid_geld.total_minted();
 
         // Trigger raid funds minting
         raid_geld.raid();
 
         // New balance should be larger
         uint256 newBalance = raid_geld.balanceOf(player1);
-        uint256 newTotalMinted = raid_geld.total_minted();
         Player memory player = raid_geld.getPlayer(player1);
         uint256 last_raided_at = player.last_raided_at;
-
         assertLt(balance, newBalance);
-        assertLt(total_minted, newTotalMinted);
 
         // Expect fail if we raid again, we need to wait a bit
         vm.expectRevert();
@@ -132,10 +123,8 @@ contract raid_geldTest is Test {
 
         // Balance should reflect that
         uint256 newestBalance = raid_geld.balanceOf(player1);
-        uint256 newestTotalMinted = raid_geld.total_minted();
         player = raid_geld.getPlayer(player1);
         uint256 last_raided_at_2 = player.last_raided_at;
-        assertLt(newTotalMinted, newestTotalMinted);
         assertLt(newBalance, newestBalance);
         assertLt(last_raided_at, last_raided_at_2);
     }
@@ -148,4 +137,33 @@ contract raid_geldTest is Test {
         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);
+    }
 }