From 306293764cb9aee25189ef26b2571d580865a10f Mon Sep 17 00:00:00 2001
From: Louise Poole <louise@datarevenue.com>
Date: Mon, 15 Apr 2024 16:56:19 +0100
Subject: [PATCH] Fix tests

---
 evm/src/balancer-v2/BalancerV2SwapAdapter.sol |  2 +-
 evm/test/BalancerV2SwapAdapter.t.sol          | 20 ++++++++++---------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/evm/src/balancer-v2/BalancerV2SwapAdapter.sol b/evm/src/balancer-v2/BalancerV2SwapAdapter.sol
index 2b5f149..040aba9 100644
--- a/evm/src/balancer-v2/BalancerV2SwapAdapter.sol
+++ b/evm/src/balancer-v2/BalancerV2SwapAdapter.sol
@@ -187,7 +187,7 @@ contract BalancerV2SwapAdapter is ISwapAdapter {
         override
         returns (Capability[] memory capabilities)
     {
-        capabilities = new Capability[](2);
+        capabilities = new Capability[](3);
         capabilities[0] = Capability.SellOrder;
         capabilities[1] = Capability.BuyOrder;
         capabilities[2] = Capability.PriceFunction;
diff --git a/evm/test/BalancerV2SwapAdapter.t.sol b/evm/test/BalancerV2SwapAdapter.t.sol
index 06ccfe0..a336af8 100644
--- a/evm/test/BalancerV2SwapAdapter.t.sol
+++ b/evm/test/BalancerV2SwapAdapter.t.sol
@@ -40,14 +40,15 @@ contract BalancerV2SwapAdapterTest is Test, ISwapAdapterTypes {
 
     function testPrice() public {
         uint256[] memory amounts = new uint256[](2);
-        amounts[0] = 100;
-        amounts[1] = 200;
-        vm.expectRevert(
-            abi.encodeWithSelector(
-                NotImplemented.selector, "BalancerV2SwapAdapter.price"
-            )
-        );
-        adapter.price(B_80BAL_20WETH_POOL_ID, BAL, WETH, amounts);
+        amounts[0] = amount0;
+        amounts[1] = amount1;
+
+        Fraction[] memory prices = adapter.price(B_80BAL_20WETH_POOL_ID, BAL, WETH, amounts);
+
+        for (uint256 i = 0; i < prices.length; i++) {
+            assertGt(prices[i].numerator, 0);
+            assertGt(prices[i].denominator, 0);
+        }
     }
 
     function testPriceSingleFuzz() public {
@@ -206,9 +207,10 @@ contract BalancerV2SwapAdapterTest is Test, ISwapAdapterTypes {
     {
         Capability[] memory res = adapter.getCapabilities(pool, t0, t1);
 
-        assertEq(res.length, 2);
+        assertEq(res.length, 3);
         assertEq(uint256(res[0]), uint256(Capability.SellOrder));
         assertEq(uint256(res[1]), uint256(Capability.BuyOrder));
+        assertEq(uint256(res[2]), uint256(Capability.PriceFunction));
     }
 
     function testGetTokens() public {