Adding fuzz flow result uncount functionality

wake/testing/fuzzing/__init__.py

@@ -1,4 +1,4 @@
-from .fuzz_test import FuzzTest, flow, invariant
+from .fuzz_test import FuzzTest, flow, invariant, FlowResult
 from .generators import (
     random_account,
     random_address,

wake/testing/fuzzing/fuzz_test.py

@@ -10,6 +10,11 @@
 from ..core import get_connected_chains
 from .generators import generate
 
+from enum import Enum
+
+class FlowResult(Enum):
+    UNCOUNT = 1
+    REPEAT = 2
 
 def flow(
     *,
@@ -82,7 +87,8 @@ def run(
             self._sequence_num = i
             self.pre_sequence()
 
-            for j in range(flows_count):
+            j = 0
+            while j <flows_count:
                 valid_flows = [
                     f
                     for f in flows
@@ -113,30 +119,38 @@ def run(
                         f"Could not find a valid flow to run.\nFlows that have reached their max_times: {max_times_flows}\nFlows that do not satisfy their precondition: {precondition_flows}"
                     )
                 flow = random.choices(valid_flows, weights=weights)[0]
-                flow_params = [
-                    generate(v)
-                    for k, v in get_type_hints(flow, include_extras=True).items()
-                    if k != "return"
-                ]
-
                 self._flow_num = j
                 self.pre_flow(flow)
-                flow(self, *flow_params)
-                flows_counter[flow] += 1
-                self.post_flow(flow)
-
-                if not dry_run:
-                    self.pre_invariants()
-                    for inv in invariants:
-                        if invariant_periods[inv] == 0:
-                            self.pre_invariant(inv)
-                            inv(self)
-                            self.post_invariant(inv)
-
-                        invariant_periods[inv] += 1
-                        if invariant_periods[inv] == getattr(inv, "period"):
-                            invariant_periods[inv] = 0
-                    self.post_invariants()
+                keep_loop = True
+                while keep_loop:
+                    flow_params = [
+                        generate(v)
+                        for k, v in get_type_hints(flow, include_extras=True).items()
+                        if k != "return"
+                    ]
+                    ret = flow(self, *flow_params)
+                    if ret == FlowResult.UNCOUNT:
+                        keep_loop = False
+                    elif ret == FlowResult.REPEAT:
+                        keep_loop = True
+                    else:
+                        j+=1
+                        flows_counter[flow] += 1
+                        self.post_flow(flow)
+                        keep_loop = False
+
+                    if not dry_run:
+                        self.pre_invariants()
+                        for inv in invariants:
+                            if invariant_periods[inv] == 0:
+                                self.pre_invariant(inv)
+                                inv(self)
+                                self.post_invariant(inv)
+
+                            invariant_periods[inv] += 1
+                            if invariant_periods[inv] == getattr(inv, "period"):
+                                invariant_periods[inv] = 0
+                        self.post_invariants()
 
             self.post_sequence()