Apply PR suggestions

- Move enum - get tacoChildApplication from Coordinator - infractions -> infractionsForRitual and value bool -> int256 - tests readability - replace `if` with `require`
nucypher · Aug 20, 2024 · c15c596 · c15c596
1 parent 2c6d97a
commit c15c596
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 48 deletions.
diff --git a/contracts/contracts/coordination/InfractionCollector.sol b/contracts/contracts/coordination/InfractionCollector.sol
@@ -12,27 +12,24 @@ contract InfractionCollector is OwnableUpgradeable {
         address indexed stakingProvider,
         InfractionType infractionType
     );
-    Coordinator public immutable coordinator;
-
-    // Reference to the TACoChildApplication contract
-    ITACoChildApplication public immutable tacoChildApplication;
-
     // infraction types
     enum InfractionType {
         MISSING_TRANSCRIPT
     }
-
+    Coordinator public immutable coordinator;
+    // Reference to the TACoChildApplication contract
+    ITACoChildApplication public immutable tacoChildApplication;
     // Mapping to keep track of reported infractions
-    mapping(uint32 ritualId => mapping(address stakingProvider => mapping(InfractionType => bool)))
-        public infractions;
+    mapping(uint32 ritualId => mapping(address stakingProvider => mapping(InfractionType => uint256)))
+        public infractionsForRitual;
 
-    constructor(Coordinator _coordinator, ITACoChildApplication _tacoChildApplication) {
+    constructor(Coordinator _coordinator) {
         require(
-            address(_coordinator) != address(0) && address(_tacoChildApplication) != address(0),
+            address(_coordinator) != address(0),
             "Contracts must be specified"
         );
         coordinator = _coordinator;
-        tacoChildApplication = _tacoChildApplication;
+        tacoChildApplication = coordinator.application();
         _disableInitializers();
     }
 
@@ -53,26 +50,22 @@ contract InfractionCollector is OwnableUpgradeable {
         for (uint256 i = 0; i < stakingProviders.length; i++) {
             // Check if the infraction has already been reported
             require(
-                !infractions[ritualId][stakingProviders[i]][InfractionType.MISSING_TRANSCRIPT],
+                infractionsForRitual[ritualId][stakingProviders[i]][InfractionType.MISSING_TRANSCRIPT] == 0,
                 "Infraction already reported"
             );
             Coordinator.Participant memory participant = coordinator.getParticipantFromProvider(
                 ritualId,
                 stakingProviders[i]
             );
-            if (participant.transcript.length == 0) {
-                // Transcript TX wasn't posted
-                // Penalize the staking provider
-                // tacoChildApplication.penalize(stakingProviders[i]);
-                infractions[ritualId][stakingProviders[i]][
-                    InfractionType.MISSING_TRANSCRIPT
-                ] = true;
-                emit InfractionReported(
-                    ritualId,
-                    stakingProviders[i],
-                    InfractionType.MISSING_TRANSCRIPT
-                );
-            }
+            require(participant.transcript.length == 0, "Transcript is not missing");
+            infractionsForRitual[ritualId][stakingProviders[i]][
+                InfractionType.MISSING_TRANSCRIPT
+            ] = 1;
+            emit InfractionReported(
+                ritualId,
+                stakingProviders[i],
+                InfractionType.MISSING_TRANSCRIPT
+            );
         }
     }
 }
diff --git a/deployment/constructor_params/lynx/infraction.yml b/deployment/constructor_params/lynx/infraction.yml
@@ -9,11 +9,9 @@ artifacts:
 constants:
   # See deployment/artifacts/lynx.json
   COORDINATOR_PROXY: "0xE9e94499bB0f67b9DBD75506ec1735486DE57770"
-  TACO_CHILD_PROXY: "0x42F30AEc1A36995eEFaf9536Eb62BD751F982D32"
 
 contracts:
   - InfractionCollector:
       proxy:
       constructor:
         _coordinator: $COORDINATOR_PROXY
-        _tacoChildApplication: $TACO_CHILD_PROXY
diff --git a/deployment/constructor_params/tapir/infraction.yml b/deployment/constructor_params/tapir/infraction.yml
@@ -9,11 +9,9 @@ artifacts:
 constants:
   # See deployment/artifacts/tapir.json
   COORDINATOR_PROXY: "0xE690b6bCC0616Dc5294fF84ff4e00335cA52C388"
-  TACO_CHILD_PROXY: "0x489287Ed5BdF7a35fEE411FBdCc47331093D0769"
 
 contracts:
   - InfractionCollector:
       proxy:
       constructor:
         _coordinator: $COORDINATOR_PROXY
-        _tacoChildApplication: $TACO_CHILD_PROXY
diff --git a/tests/test_infraction.py b/tests/test_infraction.py
@@ -11,16 +11,12 @@
 DURATION = 48 * 60 * 60
 ONE_DAY = 24 * 60 * 60
 
-RitualState = IntEnum(
-    "RitualState",
+RITUAL_ID = 0
+
+infraction_types = IntEnum(
+    "InfractionType",
     [
-        "NON_INITIATED",
-        "DKG_AWAITING_TRANSCRIPTS",
-        "DKG_AWAITING_AGGREGATIONS",
-        "DKG_TIMEOUT",
-        "DKG_INVALID",
-        "ACTIVE",
-        "EXPIRED",
+        "MISSING_TRANSCRIPT",
     ],
     start=0,
 )
@@ -112,8 +108,8 @@ def global_allow_list(project, deployer, coordinator):
 
 
 @pytest.fixture
-def infraction_collector(project, deployer, coordinator, application):
-    contract = project.InfractionCollector.deploy(coordinator.address, application.address, sender=deployer)
+def infraction_collector(project, deployer, coordinator):
+    contract = project.InfractionCollector.deploy(coordinator.address, sender=deployer)
     return contract
 
 def test_no_infractions(erc20, nodes, initiator, global_allow_list, infraction_collector, coordinator):
@@ -144,14 +140,14 @@ def test_partial_infractions(erc20, nodes, initiator, global_allow_list, infract
     transcript = os.urandom(transcript_size(len(nodes), len(nodes)))
     # post transcript for half of nodes
     for node in nodes[:len(nodes) // 2]:
-        coordinator.postTranscript(0, transcript, sender=node)
+        coordinator.postTranscript(RITUAL_ID, transcript, sender=node)
     chain.pending_timestamp += TIMEOUT * 2
-    infraction_collector.reportMissingTranscript(0, nodes, sender=initiator)
+    infraction_collector.reportMissingTranscript(RITUAL_ID, nodes[len(nodes) // 2:], sender=initiator)
     # first half of nodes should be fine, second half should be infracted
     for node in nodes[:len(nodes) // 2]:
-        assert infraction_collector.infractions(0, node, 0) == False
+        assert not infraction_collector.infractionsForRitual(RITUAL_ID, node, infraction_types.MISSING_TRANSCRIPT)
     for node in nodes[len(nodes) // 2:]:
-        assert infraction_collector.infractions(0, node, 0) == True
+        assert infraction_collector.infractionsForRitual(RITUAL_ID, node, infraction_types.MISSING_TRANSCRIPT)
 
 def test_report_infractions(erc20, nodes, initiator, global_allow_list, infraction_collector, coordinator, chain):
     cost = coordinator.getRitualInitiationCost(nodes, DURATION)
@@ -163,9 +159,9 @@ def test_report_infractions(erc20, nodes, initiator, global_allow_list, infracti
         nodes, initiator, DURATION, global_allow_list.address, sender=initiator
     )
     chain.pending_timestamp += TIMEOUT * 2
-    infraction_collector.reportMissingTranscript(0, nodes, sender=initiator)
+    infraction_collector.reportMissingTranscript(RITUAL_ID, nodes, sender=initiator)
     for node in nodes:
-        assert infraction_collector.infractions(0, node, 0) == True
+        assert infraction_collector.infractionsForRitual(RITUAL_ID, node, infraction_types.MISSING_TRANSCRIPT)
 
 def test_cant_report_infractions_twice(erc20, nodes, initiator, global_allow_list, infraction_collector, coordinator, chain):
     cost = coordinator.getRitualInitiationCost(nodes, DURATION)
@@ -177,7 +173,7 @@ def test_cant_report_infractions_twice(erc20, nodes, initiator, global_allow_lis
         nodes, initiator, DURATION, global_allow_list.address, sender=initiator
     )
     chain.pending_timestamp += TIMEOUT * 2
-    infraction_collector.reportMissingTranscript(0, nodes, sender=initiator)
+    infraction_collector.reportMissingTranscript(RITUAL_ID, nodes, sender=initiator)
 
     with ape.reverts("Infraction already reported"):
-        infraction_collector.reportMissingTranscript(0, nodes, sender=initiator)
+        infraction_collector.reportMissingTranscript(RITUAL_ID, nodes, sender=initiator)