added more tests + ref

Makefile

@@ -28,13 +28,17 @@ update:
 build:
 	forge clean && forge build
 
-# Generate gas snapshot under snapshots directory
+# Generate gas snapshot
 snapshot:
 	forge snapshot
 
-# Test the contracts on forked base-sepolia network
+# Test the contracts forked base-sepolia network with 4 parallel jobs
 test:
-	forge clean && forge test --fork-url $(BASE_TEST_RPC_URL)
+	forge clean && forge test --fork-url $(BASE_TEST_RPC_URL) --no-match-contract "InvariantTest" --jobs 4
+
+# Run invariant tests on local network with 4 parallel jobs
+test-inv:
+	forge clean && forge test --match-contract "InvariantTest" --jobs 4
 
 anvil:
 	anvil --fork-url $(BASE_TEST_RPC_URL)
@@ -57,7 +61,7 @@ fmt:
 
 # Coverage
 cov:
-	forge coverage --no-match-coverage "(test|mock|script)"
+	forge clean && forge coverage --no-match-coverage "(test|mock|script)" --no-match-contract "InvariantTest" --jobs 4
 
 # Verify contract on blockscout
 verify:

README.md

@@ -14,12 +14,17 @@ make build
 
 ## Test
 
-Run tests on forked base-sepolia:
+Run tests on forked base-sepolia with:
 
 ```sh
 make test
 ```
 
+Run invariant tests on local with:
+```sh
+make test-inv
+```
+
 ## Deployment
 
 **Step 1.**

foundry.toml

@@ -5,11 +5,23 @@ test = 'test'
 script = 'script'
 out = 'out'
 cache_path  = 'cache'
+extra_output = ['storageLayout']
+
 ffi = true
 ast = true
 build_info = true
 optimizer = true
-extra_output = ['storageLayout']
+
+# fuzzing options
+fuzz_runs = 100
+
+# invariant options
+invariant_runs = 20
+
+# block timestamp starts from 10
+block_timestamp = 10
+
+# fs permissions for deployment
 fs_permissions = [{ access = "read", path = "out" }, { access = "write", path = "deployment" }]
 remappings = [
     "@openzeppelin/contracts-upgradeable/=lib/openzeppelin-contracts-upgradeable/contracts/",
@@ -18,4 +30,5 @@ remappings = [
     "@firstbatch/dria-oracle-contracts/=lib/dria-oracle-contracts/src/"
 ]
 
+
 # See more config options https://github.com/foundry-rs/foundry/blob/master/crates/config/README.md#all-options

script/Deploy.s.sol

@@ -53,7 +53,10 @@ contract Deploy is Script {
         // deploy llm contracts
         address registryProxy = Upgrades.deployUUPSProxy(
             "LLMOracleRegistry.sol",
-            abi.encodeCall(LLMOracleRegistry.initialize, (genStake, valStake, address(config.token())))
+            abi.encodeCall(
+                LLMOracleRegistry.initialize,
+                (genStake, valStake, address(config.token()), config.minRegistrationTime())
+            )
         );
 
         // wrap proxy with the LLMOracleRegistry
@@ -65,7 +68,15 @@ contract Deploy is Script {
             "LLMOracleCoordinator.sol",
             abi.encodeCall(
                 LLMOracleCoordinator.initialize,
-                (address(oracleRegistry), address(config.token()), platformFee, genFee, valFee)
+                (
+                    address(oracleRegistry),
+                    address(config.token()),
+                    platformFee,
+                    genFee,
+                    valFee,
+                    config.minScore(),
+                    config.maxScore()
+                )
             )
         );
 
@@ -87,6 +98,7 @@ contract Deploy is Script {
             uint256 platformFee,
             uint256 maxAssetCount,
             uint256 minAssetPrice,
+            /* timestamp */
             ,
             uint8 maxBuyerAgentFee
         ) = config.marketParams();

script/HelperConfig.s.sol

@@ -13,8 +13,8 @@ struct Stakes {
 
 struct Fees {
     uint256 platformFee;
-    uint256 generatorFee;
-    uint256 validatorFee;
+    uint256 generationFee;
+    uint256 validationFee;
 }
 
 contract HelperConfig is Script {
@@ -25,10 +25,14 @@ contract HelperConfig is Script {
     Fees public fees;
     WETH9 public token;
 
+    uint256 public minRegistrationTime; // in seconds
+    uint256 public minScore;
+    uint256 public maxScore;
+
     constructor() {
         // set deployment parameters
         stakes = Stakes({generatorStakeAmount: 0.0001 ether, validatorStakeAmount: 0.000001 ether});
-        fees = Fees({platformFee: 0.0001 ether, generatorFee: 0.0001 ether, validatorFee: 0.0001 ether});
+        fees = Fees({platformFee: 0.0001 ether, generationFee: 0.0001 ether, validationFee: 0.0001 ether});
         taskParams = LLMOracleTaskParameters({difficulty: 2, numGenerations: 1, numValidations: 1});
 
         marketParams = SwanMarketParameters({
@@ -42,6 +46,10 @@ contract HelperConfig is Script {
             maxBuyerAgentFee: 75 // percentage
         });
 
+        minRegistrationTime = 1 days;
+        maxScore = type(uint8).max; // 255
+        minScore = 1;
+
         // for base sepolia
         if (block.chainid == 84532) {
             // use deployed weth

src/BuyerAgent.sol

@@ -47,6 +47,22 @@ contract BuyerAgent is Ownable {
     /// @notice The task was already processed, via `purchase` or `updateState`.
     error TaskAlreadyProcessed();
 
+    /*//////////////////////////////////////////////////////////////
+                                 EVENTS
+    //////////////////////////////////////////////////////////////*/
+
+    /// @notice Emitted when a state update request is made.
+    event StateRequest(uint256 indexed taskId, uint256 indexed round);
+
+    /// @notice Emitted when a purchase request is made.
+    event PurchaseRequest(uint256 indexed taskId, uint256 indexed round);
+
+    /// @notice Emitted when a purchase is made.
+    event Purchase(uint256 indexed round);
+
+    /// @notice Emitted when the state is updated.
+    event StateUpdate(uint256 indexed round);
+
     /*//////////////////////////////////////////////////////////////
                                  STORAGE
     //////////////////////////////////////////////////////////////*/
@@ -178,6 +194,8 @@ contract BuyerAgent is Ownable {
 
         oracleStateRequests[round] =
             swan.coordinator().request(SwanBuyerStateOracleProtocol, _input, _models, swan.getOracleParameters());
+
+        emit StateRequest(oracleStateRequests[round], round);
     }
 
     /// @notice Calls the LLMOracleCoordinator & pays for the oracle fees to make a purchase request.
@@ -193,6 +211,8 @@ contract BuyerAgent is Ownable {
 
         oraclePurchaseRequests[round] =
             swan.coordinator().request(SwanBuyerPurchaseOracleProtocol, _input, _models, swan.getOracleParameters());
+
+        emit PurchaseRequest(oraclePurchaseRequests[round], round);
     }
 
     /// @notice Function to update the Buyer state.
@@ -214,6 +234,8 @@ contract BuyerAgent is Ownable {
 
         // update taskId as completed
         isOracleRequestProcessed[taskId] = true;
+
+        emit StateUpdate(round);
     }
 
     /// @notice Function to buy the asset from the Swan with the given assed address.
@@ -254,6 +276,8 @@ contract BuyerAgent is Ownable {
 
         // update taskId as completed
         isOracleRequestProcessed[taskId] = true;
+
+        emit Purchase(round);
     }
 
     /// @notice Function to withdraw the tokens from the contract.

src/Swan.sol

@@ -4,18 +4,21 @@ pragma solidity ^0.8.20;
 import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
-import {LLMOracleCoordinator} from "@firstbatch/dria-oracle-contracts/LLMOracleCoordinator.sol";
-import {LLMOracleTaskParameters} from "@firstbatch/dria-oracle-contracts/LLMOracleTask.sol";
+import {LLMOracleCoordinator, LLMOracleTaskParameters} from "@firstbatch/dria-oracle-contracts/LLMOracleCoordinator.sol";
 import {BuyerAgentFactory, BuyerAgent} from "./BuyerAgent.sol";
 import {SwanAssetFactory, SwanAsset} from "./SwanAsset.sol";
 import {SwanManager, SwanMarketParameters} from "./SwanManager.sol";
+import {Math} from "@openzeppelin/contracts/utils/math/Math.sol";
 
-// Protocol strings for Swan, checked in the Oracle.
+// @dev Protocol strings for Swan, checked in the Oracle.
 bytes32 constant SwanBuyerPurchaseOracleProtocol = "swan-buyer-purchase/0.1.0";
 bytes32 constant SwanBuyerStateOracleProtocol = "swan-buyer-state/0.1.0";
+
+/// @dev Used to calculate the fee for the buyer agent to be able to compute correct amount.
 uint256 constant BASIS_POINTS = 10_000;
 
 contract Swan is SwanManager, UUPSUpgradeable {
+    using Math for uint256;
     /*//////////////////////////////////////////////////////////////
                                  ERRORS
     //////////////////////////////////////////////////////////////*/
@@ -307,15 +310,16 @@ contract Swan is SwanManager, UUPSUpgradeable {
     /// @notice Function to transfer the royalties to the seller & Dria.
     function transferRoyalties(AssetListing storage asset) internal {
         // calculate fees
-        uint256 buyerFee = (asset.price * asset.feeRoyalty) / (BASIS_POINTS * BASIS_POINTS);
-        uint256 driaFee = (buyerFee * asset.feeRoyalty * getCurrentMarketParameters().platformFee) / (BASIS_POINTS * BASIS_POINTS);
+        uint256 totalFee = Math.mulDiv(asset.price, (asset.feeRoyalty * 100), BASIS_POINTS);
+        uint256 driaFee = Math.mulDiv(totalFee, (getCurrentMarketParameters().platformFee * 100), BASIS_POINTS);
+        uint256 buyerFee = totalFee - driaFee;
 
         // first, Swan receives the entire fee from seller
         // this allows only one approval from the seller's side
-        token.transferFrom(asset.seller, address(this), buyerFee);
+        token.transferFrom(asset.seller, address(this), totalFee);
 
         // send the buyer's portion to them
-        token.transfer(asset.buyer, buyerFee - driaFee);
+        token.transfer(asset.buyer, buyerFee);
 
         // then it sends the remaining to Swan owner
         token.transfer(owner(), driaFee);

test/Helper.t.sol

@@ -65,21 +65,25 @@ abstract contract Helper is Test {
     uint256 amountPerRound = 0.015 ether;
     uint8 feeRoyalty = 2;
 
-    // Default scores for validation
+    /// @dev Default scores for validation
     uint256[] scores = [1, 5, 70];
 
+    uint256 public minRegistrationTime = 1 days; // in seconds
+    uint256 public minScore = 1;
+    uint256 public maxScore = type(uint8).max; // 255
+
     /// @notice The given nonce is not a valid proof-of-work.
     error InvalidNonceFromHelperTest(uint256 taskId, uint256 nonce, uint256 computedNonce, address caller);
 
-    // Set parameters for the test
+    // @dev Set parameters for the test
     function setUp() public deployment {
         dria = vm.addr(1);
         validators = [vm.addr(2), vm.addr(3), vm.addr(4)];
         generators = [vm.addr(5), vm.addr(6), vm.addr(7)];
         buyerAgentOwners = [vm.addr(8), vm.addr(9)];
         sellers = [vm.addr(10), vm.addr(11)];
 
-        oracleParameters = LLMOracleTaskParameters({difficulty: 1, numGenerations: 1, numValidations: 1});
+        oracleParameters = LLMOracleTaskParameters({difficulty: 1, numGenerations: 2, numValidations: 1});
         marketParameters = SwanMarketParameters({
             withdrawInterval: 300, // 5 minutes
             sellInterval: 360,
@@ -92,7 +96,7 @@ abstract contract Helper is Test {
         });
 
         stakes = Stakes({generatorStakeAmount: 0.01 ether, validatorStakeAmount: 0.01 ether});
-        fees = Fees({platformFee: 0.0001 ether, generatorFee: 0.0002 ether, validatorFee: 0.00003 ether});
+        fees = Fees({platformFee: 1, generationFee: 0.0002 ether, validationFee: 0.00003 ether});
 
         for (uint96 i = 0; i < buyerAgentOwners.length; i++) {
             buyerAgentParameters.push(
@@ -125,7 +129,8 @@ abstract contract Helper is Test {
         address registryProxy = Upgrades.deployUUPSProxy(
             "LLMOracleRegistry.sol",
             abi.encodeCall(
-                LLMOracleRegistry.initialize, (stakes.generatorStakeAmount, stakes.validatorStakeAmount, address(token))
+                LLMOracleRegistry.initialize,
+                (stakes.generatorStakeAmount, stakes.validatorStakeAmount, address(token), minRegistrationTime)
             )
         );
         oracleRegistry = LLMOracleRegistry(registryProxy);
@@ -134,7 +139,15 @@ abstract contract Helper is Test {
             "LLMOracleCoordinator.sol",
             abi.encodeCall(
                 LLMOracleCoordinator.initialize,
-                (address(oracleRegistry), address(token), fees.platformFee, fees.generatorFee, fees.validatorFee)
+                (
+                    address(oracleRegistry),
+                    address(token),
+                    fees.platformFee,
+                    fees.generationFee,
+                    fees.validationFee,
+                    minScore,
+                    maxScore
+                )
             )
         );
         oracleCoordinator = LLMOracleCoordinator(coordinatorProxy);
@@ -169,6 +182,17 @@ abstract contract Helper is Test {
         vm.label(address(swanAssetFactory), "SwanAssetFactory");
     }
 
+    /// @notice Add validators to the whitelist.
+    modifier addValidatorsToWhitelist() {
+        vm.prank(dria);
+        oracleRegistry.addToWhitelist(validators);
+
+        for (uint256 i; i < validators.length; i++) {
+            vm.assertTrue(oracleRegistry.whitelisted(validators[i]));
+        }
+        _;
+    }
+
     /// @notice Register generators and validators
     modifier registerOracles() {
         for (uint256 i = 0; i < generators.length; i++) {
@@ -460,11 +484,12 @@ abstract contract Helper is Test {
 
         // respond
         safeRespond(generators[0], encodedOutput, taskId);
+        safeRespond(generators[1], encodedOutput, taskId);
 
         // validate
         safeValidate(validators[0], taskId);
 
-        assert(token.balanceOf(address(buyerAgent)) > assetPrice);
+        assertGe(token.balanceOf(address(buyerAgent)), assetPrice);
 
         // purchase and check event logs
         vm.recordLogs();