From 341393ad4698c9e810f7c655afe28a6e760b8054 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stas=20SU=C8=98COV?= <stas@nerd.ro>
Date: Wed, 13 Nov 2024 21:11:49 +0000
Subject: [PATCH] RewardsSugar: added forRoot helper.

---
 contracts/RewardsSugar.vy   | 20 ++++++++++++++++++++
 env.optimism                |  2 +-
 tests/test_rewards_sugar.py | 15 +++++++++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/contracts/RewardsSugar.vy b/contracts/RewardsSugar.vy
index 51cd4cd..0916590 100644
--- a/contracts/RewardsSugar.vy
+++ b/contracts/RewardsSugar.vy
@@ -347,3 +347,23 @@ def _pool_rewards(_venft_id: uint256, _pool: address, _gauge: address) \
     )
 
   return col
+
+
+@external
+@view
+def forRoot(_root_pool: address) -> address[3]:
+  """
+  @notice Returns rewards addresses for the root pool
+  @param _root_pool the root pool address to map to
+  @return Array with the root gauge, fee and incentive addresses
+  """
+  if chain.id not in lp_shared.ROOT_CHAIN_IDS:
+    return empty(address[3])
+
+  gauge: address = staticcall lp_shared.voter.gauges(_root_pool)
+
+  return [
+    gauge,
+    staticcall lp_shared.voter.gaugeToFees(gauge),
+    staticcall lp_shared.voter.gaugeToBribe(gauge)
+  ]
diff --git a/env.optimism b/env.optimism
index 81e720c..7649b1a 100644
--- a/env.optimism
+++ b/env.optimism
@@ -14,6 +14,6 @@ TEST_ADDRESS_10=0x892ff98a46e5bd141e2d12618f4b2fe6284debac
 TEST_ALM_ADDRESS_10=0x892ff98a46e5bd141e2d12618f4b2fe6284debac
 
 LP_SUGAR_ADDRESS_10=0x5B3B370C85c8816cBB4BCE5f647dCa3e06c421d7
-REWARDS_SUGAR_ADDRESS_10=0x30A6cA7305Cb8a53A530533a680A06f5cB7f11Fd
+REWARDS_SUGAR_ADDRESS_10=0x62CCFB2496f49A80B0184AD720379B529E9152fB
 VE_SUGAR_ADDRESS_10=0x94f913362b232e31daB49a1aFB775cfd25DaA6a1
 RELAY_SUGAR_ADDRESS_10=0xb8307e5842B9aeE75C704183F0355076aa74b4e2
diff --git a/tests/test_rewards_sugar.py b/tests/test_rewards_sugar.py
index 4d3705e..669fa37 100644
--- a/tests/test_rewards_sugar.py
+++ b/tests/test_rewards_sugar.py
@@ -110,3 +110,18 @@ def test_epochsLatest_limit_offset(
 
     assert lepoch.lp == pepoch.lp
     assert lepoch.ts == pepoch.ts
+
+
+@pytest.mark.skipif(int(CHAIN_ID) not in [10], reason="Only Optimism")
+def test_forRoot(sugar_contract, lp_sugar_contract, LpStruct):
+    first_lp = LpStruct(*lp_sugar_contract.byIndex(1))
+
+    # Use `lp` instead of `root` for testing
+    addresses = sugar_contract.forRoot(first_lp.lp)
+
+    assert addresses is not None
+
+    assert len(addresses) == 3
+    assert addresses[0] == first_lp.gauge
+    assert addresses[1] == first_lp.fee
+    assert addresses[2] == first_lp.bribe