random.rs

//! Contains various tests for `random*` and `fresh*` cheatcodes.

use alloy_primitives::U256;
use foundry_config::{Config, FuzzConfig};
use foundry_test_utils::{forgetest_init, str, util::OutputExt};

// tests that `forge test` with a seed produces deterministic random values for uint and addresses.
forgetest_init!(deterministic_randomness_with_seed, |prj, cmd| {
    prj.wipe_contracts();
    prj.add_test(
        "DeterministicRandomnessTest.t.sol",
        r#"
import {Test, console} from "forge-std/Test.sol";

contract DeterministicRandomnessTest is Test {

    function testDeterministicRandomUint() public {
        console.log(vm.randomUint());
        console.log(vm.randomUint());
        console.log(vm.randomUint());
    }

    function testDeterministicRandomUintRange() public {
        uint256 min = 0;
        uint256 max = 1000000000;
        console.log(vm.randomUint(min, max));
        console.log(vm.randomUint(min, max));
        console.log(vm.randomUint(min, max));
    }

    function testDeterministicRandomAddress() public {
        console.log(vm.randomAddress());
        console.log(vm.randomAddress());
        console.log(vm.randomAddress());
    }
}
"#,
    )
    .unwrap();

    // Extracts the test result section from the DeterministicRandomnessTest contract output.
    fn extract_test_result(out: &str) -> &str {
        let start = out
            .find("for test/DeterministicRandomnessTest.t.sol:DeterministicRandomnessTest")
            .unwrap();
        let end = out.find("Suite result: ok.").unwrap();
        &out[start..end]
    }

    // Run the test twice with the same seed and verify the outputs are the same.
    let seed1 = "0xa1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2";
    let out1 = cmd
        .args(["test", "--fuzz-seed", seed1, "-vv"])
        .assert_success()
        .get_output()
        .stdout_lossy();
    let res1 = extract_test_result(&out1);

    let out2 = cmd
        .forge_fuse()
        .args(["test", "--fuzz-seed", seed1, "-vv"])
        .assert_success()
        .get_output()
        .stdout_lossy();
    let res2 = extract_test_result(&out2);

    assert_eq!(res1, res2);

    // Run the test with another seed and verify the output differs.
    let seed2 = "0xb1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2";
    let out3 = cmd
        .forge_fuse()
        .args(["test", "--fuzz-seed", seed2, "-vv"])
        .assert_success()
        .get_output()
        .stdout_lossy();
    let res3 = extract_test_result(&out3);
    assert_ne!(res3, res1);

    // Run the test without a seed and verify the outputs differs once again.
    cmd.forge_fuse();
    let out4 = cmd.args(["test", "-vv"]).assert_success().get_output().stdout_lossy();
    let res4 = extract_test_result(&out4);
    assert_ne!(res4, res1);
    assert_ne!(res4, res3);
});

// tests `freshStorage` cheatcode.
forgetest_init!(test_fresh_storage, |prj, cmd| {
    prj.wipe_contracts();
    prj.insert_ds_test();
    prj.insert_vm();
    prj.clear();

    prj.add_source(
        "Counter.t.sol",
        r#"pragma solidity 0.8.24;
import {Vm} from "./Vm.sol";
import {DSTest} from "./test.sol";
contract Counter {
    uint256 public a;
    address public b;
    int8 public c;
    address[] public owners;

    function setA(uint256 _a) public {
        a = _a;
    }

    function setB(address _b) public {
        b = _b;
    }

    function getOwner(uint256 pos) public view returns (address) {
        return owners[pos];
    }

    function setOwner(uint256 pos, address owner) public {
        owners[pos] = owner;
    }
}

contract CounterFreshStorageTest is DSTest {
    Vm vm = Vm(HEVM_ADDRESS);

    function test_fresh_storage() public {
        uint256 index = 55;
        Counter counter = new Counter();
        vm.freshStorage(address(counter));

        // Next call would fail with array out of bounds without freshStorage
        address owner = counter.getOwner(index);
        // Subsequent calls should retrieve same value
        assertEq(counter.getOwner(index), owner);
        // Change slot and make sure new value retrieved
        counter.setOwner(index, address(111));
        assertEq(counter.getOwner(index), address(111));
    }


    function test_fresh_storage_warm() public {
        Counter counter = new Counter();
        vm.freshStorage(address(counter));
        assertGt(counter.a(), 0);
        counter.setA(0);
        // This should remain 0 if explicitly set.
        assertEq(counter.a(), 0);
        counter.setA(11);
        assertEq(counter.a(), 11);
    }

    function test_fresh_storage_multiple_read_writes() public {
        Counter counter = new Counter();
        vm.freshStorage(address(counter));
        uint256 slot1 = vm.randomUint(0, 100);
        uint256 slot2 = vm.randomUint(0, 100);
        require(slot1 != slot2, "random positions should be different");
        address alice = counter.owners(slot1);
        address bob = counter.owners(slot2);
        require(alice != bob, "random storage values should be different");
        counter.setOwner(slot1, bob);
        counter.setOwner(slot2, alice);
        assertEq(alice, counter.owners(slot2));
        assertEq(bob, counter.owners(slot1));
    }
}
     "#,
    )
    .unwrap();

    cmd.args(["test"]).assert_success().stdout_eq(str![[r#"
...
[PASS] test_fresh_storage() ([GAS])
[PASS] test_fresh_storage_multiple_read_writes() ([GAS])
[PASS] test_fresh_storage_warm() ([GAS])
...
"#]]);
});

// tests deterministic `freshStorage` cheatcode.
forgetest_init!(test_fresh_storage_with_seed, |prj, cmd| {
    prj.wipe_contracts();
    prj.insert_ds_test();
    prj.insert_vm();
    prj.clear();

    let config = Config {
        fuzz: { FuzzConfig { seed: Some(U256::from(100)), ..Default::default() } },
        ..Default::default()
    };
    prj.write_config(config);

    prj.add_source(
        "Counter.t.sol",
        r#"pragma solidity 0.8.24;
import {Vm} from "./Vm.sol";
import {DSTest} from "./test.sol";
contract Counter {
    uint256[] public a;
    address[] public b;
    int8[] public c;
    bytes32[] public d;
}

contract CounterFreshStorageTest is DSTest {
    Vm vm = Vm(HEVM_ADDRESS);

    function test_fresh_storage_with_seed() public {
        Counter counter = new Counter();
        vm.freshStorage(address(counter));
        assertEq(counter.a(11), 85286582241781868037363115933978803127245343755841464083427462398552335014708);
        assertEq(counter.b(22), 0x939180Daa938F9e18Ff0E76c112D25107D358B02);
        assertEq(counter.c(33), -104);
        assertEq(counter.d(44), 0x6c178fa9c434f142df61a5355cc2b8d07be691b98dabf5b1a924f2bce97a19c7);
    }
}
     "#,
    )
    .unwrap();

    cmd.args(["test"]).assert_success().stdout_eq(str![[r#"
...
[PASS] test_fresh_storage_with_seed() ([GAS])
...
"#]]);
});

// Kontrol `symbolicStorage` cheatcode tests ported to `freshStorage`.
forgetest_init!(test_fresh_storage_kontrol, |prj, cmd| {
    prj.wipe_contracts();
    prj.insert_ds_test();
    prj.insert_vm();
    prj.clear();

    let config = Config {
        fuzz: { FuzzConfig { seed: Some(U256::from(100)), ..Default::default() } },
        ..Default::default()
    };
    prj.write_config(config);

    prj.add_source(
        "Counter.t.sol",
        r#"pragma solidity 0.8.24;
import {Vm} from "./Vm.sol";
import {DSTest} from "./test.sol";
contract SymbolicStore {
    uint256 public testNumber = 1337; // slot 0
    constructor() {}
}

contract SymbolicStorageTest is DSTest {
    Vm vm = Vm(HEVM_ADDRESS);

    function test_SymbolicStorage() public {
         uint256 slot = vm.randomUint(0, 100);
         address addr = 0xEA674fdDe714fd979de3EdF0F56AA9716B898ec8;
         vm.freshStorage(addr);
         bytes32 value = vm.load(addr, bytes32(slot));
         assertEq(uint256(value), 85286582241781868037363115933978803127245343755841464083427462398552335014708);
         // Load slot again and make sure we get same value.
         bytes32 value1 = vm.load(addr, bytes32(slot));
         assertEq(uint256(value), uint256(value1));
    }

    function test_SymbolicStorage1() public {
        uint256 slot = vm.randomUint(0, 100);
        SymbolicStore myStore = new SymbolicStore();
        vm.freshStorage(address(myStore));
        bytes32 value = vm.load(address(myStore), bytes32(uint256(slot)));
        assertEq(uint256(value), 85286582241781868037363115933978803127245343755841464083427462398552335014708);
    }

    function testEmptyInitialStorage(uint256 slot) public {
        bytes32 storage_value = vm.load(address(vm), bytes32(slot));
        assertEq(uint256(storage_value), 0);
    }
}
     "#,
    )
    .unwrap();

    cmd.args(["test"]).assert_success().stdout_eq(str![[r#"
...
[PASS] testEmptyInitialStorage(uint256) (runs: 256, [AVG_GAS])
[PASS] test_SymbolicStorage() ([GAS])
[PASS] test_SymbolicStorage1() ([GAS])
...
"#]]);
});