Skip to content

Commit

Permalink
update tests
Browse files Browse the repository at this point in the history
  • Loading branch information
caglacelik committed Dec 2, 2024
1 parent 679359a commit d766cc5
Show file tree
Hide file tree
Showing 4 changed files with 294 additions and 335 deletions.
149 changes: 38 additions & 111 deletions test/BuyerAgent.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -14,107 +14,15 @@ import {LLMOracleRegistry} from "@firstbatch/dria-oracle-contracts/LLMOracleRegi
import {Swan} from "../src/Swan.sol";

contract BuyerAgentTest is Helper {
address agentOwner;
BuyerAgent agent;

modifier deployment() {
token = new WETH9();

// deploy llm contracts
vm.startPrank(dria);

address registryProxy = Upgrades.deployUUPSProxy(
"LLMOracleRegistry.sol",
abi.encodeCall(
LLMOracleRegistry.initialize, (stakes.generatorStakeAmount, stakes.validatorStakeAmount, address(token))
)
);

oracleRegistry = LLMOracleRegistry(registryProxy);

address coordinatorProxy = Upgrades.deployUUPSProxy(
"LLMOracleCoordinator.sol",
abi.encodeCall(
LLMOracleCoordinator.initialize,
(address(oracleRegistry), address(token), fees.platformFee, fees.generatorFee, fees.validatorFee)
)
);
oracleCoordinator = LLMOracleCoordinator(coordinatorProxy);

// deploy factory contracts
buyerAgentFactory = new BuyerAgentFactory();
swanAssetFactory = new SwanAssetFactory();

require(address(buyerAgentFactory) != address(0), "BuyerAgentFactory not deployed");
require(address(swanAssetFactory) != address(0), "SwanAssetFactory not deployed");

// deploy swan
address swanProxy = Upgrades.deployUUPSProxy(
"Swan.sol",
abi.encodeCall(
Swan.initialize,
(
marketParameters,
oracleParameters,
address(oracleCoordinator),
address(token),
address(buyerAgentFactory),
address(swanAssetFactory)
)
)
);
swan = Swan(swanProxy);
vm.stopPrank();

vm.label(address(swan), "Swan");
vm.label(address(token), "WETH");
vm.label(address(this), "BuyerAgentTest");
vm.label(address(oracleRegistry), "LLMOracleRegistry");
vm.label(address(oracleCoordinator), "LLMOracleCoordinator");
vm.label(address(buyerAgentFactory), "BuyerAgentFactory");
vm.label(address(swanAssetFactory), "SwanAssetFactory");
_;
}

modifier createBuyers() override {
for (uint256 i = 0; i < buyerAgentOwners.length; i++) {
// fund buyer agent owner
deal(address(token), buyerAgentOwners[i], 3 ether);

vm.startPrank(buyerAgentOwners[i]);
BuyerAgent buyerAgent = swan.createBuyer(
buyerAgentParameters[i].name,
buyerAgentParameters[i].description,
buyerAgentParameters[i].feeRoyalty,
buyerAgentParameters[i].amountPerRound
);

buyerAgents.push(buyerAgent);
vm.label(address(buyerAgent), string.concat("BuyerAgent#", vm.toString(i + 1)));

// transfer tokens to agent
token.transfer(address(buyerAgent), amountPerRound);
assertEq(token.balanceOf(address(buyerAgent)), amountPerRound);
vm.stopPrank();
}

agentOwner = buyerAgentOwners[0];
agent = buyerAgents[0];

currPhase = BuyerAgent.Phase.Sell;
currRound = 0;
_;
}

/// @notice Buyer agent should be in sell phase
function test_InSellPhase() external deployment createBuyers {
function test_InSellPhase() external createBuyers {
// get curr phase
(, BuyerAgent.Phase _phase,) = agent.getRoundPhase();
assertEq(uint8(_phase), uint8(currPhase));
}

/// @dev Agent owner cannot set feeRoyalty in sell phase
function test_RevertWhen_SetRoyaltyInSellPhase() external deployment createBuyers {
function test_RevertWhen_SetRoyaltyInSellPhase() external createBuyers {
vm.prank(agentOwner);
vm.expectRevert(
abi.encodeWithSelector(BuyerAgent.InvalidPhase.selector, BuyerAgent.Phase.Sell, BuyerAgent.Phase.Withdraw)
Expand All @@ -123,17 +31,19 @@ contract BuyerAgentTest is Helper {
}

/// @notice Test that the buyer agent is in buy phase
function test_InBuyPhase() external deployment createBuyers {
vm.warp(agent.createdAt() + marketParameters.sellInterval);
function test_InBuyPhase() external createBuyers {
uint256 timeToBuyPhaseOfTheFirstRound = agent.createdAt() + swan.getCurrentMarketParameters().sellInterval;
increaseTime(timeToBuyPhaseOfTheFirstRound, agent, BuyerAgent.Phase.Buy, 0);

currPhase = BuyerAgent.Phase.Buy;

(, BuyerAgent.Phase _phase,) = agent.getRoundPhase();
assertEq(uint8(_phase), uint8(currPhase));
}

/// @dev Agent owner cannot set amountPerRound in buy phase
function test_RevertWhen_SetAmountPerRoundInBuyPhase() external deployment createBuyers {
vm.warp(agent.createdAt() + marketParameters.sellInterval);
function test_RevertWhen_SetAmountPerRoundInBuyPhase() external createBuyers {
increaseTime(agent.createdAt() + marketParameters.sellInterval, agent, BuyerAgent.Phase.Buy, 0);

vm.prank(agentOwner);
vm.expectRevert(
Expand All @@ -143,9 +53,9 @@ contract BuyerAgentTest is Helper {
}

/// @notice Test that the buyer agent owner cannot withdraw in buy phase
function test_RevertWhen_WithdrawInBuyPhase() external deployment createBuyers {
function test_RevertWhen_WithdrawInBuyPhase() external createBuyers {
// owner cannot withdraw more than minFundAmount from his agent
vm.warp(agent.createdAt() + marketParameters.sellInterval);
increaseTime(agent.createdAt() + marketParameters.sellInterval, agent, BuyerAgent.Phase.Buy, 0);

// get the contract balance
uint256 treasuary = agent.treasury();
Expand All @@ -157,23 +67,30 @@ contract BuyerAgentTest is Helper {
}

/// @notice Test that the non-owner cannot withdraw
function test_RevertWhen_WithdrawByAnotherOwner() external deployment createBuyers {
function test_RevertWhen_WithdrawByAnotherOwner() external createBuyers {
// feeRoyalty can be set only in withdraw phase by only agent owner
vm.warp(agent.createdAt() + marketParameters.sellInterval + marketParameters.buyInterval);
increaseTime(
agent.createdAt() + marketParameters.sellInterval + marketParameters.buyInterval,
agent,
BuyerAgent.Phase.Withdraw,
0
);
currPhase = BuyerAgent.Phase.Withdraw;

(, BuyerAgent.Phase _phase,) = agent.getRoundPhase();
assertEq(uint8(_phase), uint8(currPhase));

// not allowed to withdraw by non owner
vm.prank(buyerAgentOwners[1]);
vm.expectRevert(abi.encodeWithSelector(BuyerAgent.Unauthorized.selector, buyerAgentOwners[1]));
agent.withdraw(1 ether);
}

/// @notice Test that the buyer agent owner must set feeRoyalty between 1-100
function test_RevertWhen_SetFeeWithInvalidRoyalty() external deployment createBuyers {
vm.warp(agent.createdAt() + marketParameters.sellInterval + marketParameters.buyInterval);
function test_RevertWhen_SetFeeWithInvalidRoyalty() external createBuyers {
increaseTime(
agent.createdAt() + marketParameters.sellInterval + marketParameters.buyInterval,
agent,
BuyerAgent.Phase.Withdraw,
0
);

uint96 biggerRoyalty = 1000;
uint96 smallerRoyalty = 0;
Expand All @@ -188,8 +105,13 @@ contract BuyerAgentTest is Helper {
}

/// @notice Test that the buyer agent owner can set feeRoyalty and amountPerRound in withdraw phase
function test_SetRoyaltyAndAmountPerRound() external deployment createBuyers {
vm.warp(agent.createdAt() + marketParameters.sellInterval + marketParameters.buyInterval);
function test_SetRoyaltyAndAmountPerRound() external createBuyers {
increaseTime(
agent.createdAt() + marketParameters.sellInterval + marketParameters.buyInterval,
agent,
BuyerAgent.Phase.Withdraw,
0
);

uint96 newFeeRoyalty = 20;
uint256 newAmountPerRound = 0.25 ether;
Expand All @@ -205,8 +127,13 @@ contract BuyerAgentTest is Helper {
}

/// @notice Test that the buyer agent owner can withdraw in withdraw phase
function test_WithdrawInWithdrawPhase() external deployment createBuyers {
vm.warp(agent.createdAt() + marketParameters.sellInterval + marketParameters.buyInterval);
function test_WithdrawInWithdrawPhase() external createBuyers {
increaseTime(
agent.createdAt() + marketParameters.sellInterval + marketParameters.buyInterval,
agent,
BuyerAgent.Phase.Withdraw,
0
);

vm.startPrank(agentOwner);
agent.withdraw(uint96(token.balanceOf(address(agent))));
Expand Down
Loading

0 comments on commit d766cc5

Please sign in to comment.