From 945c089400d0cbbdd5b22a0f5885d265835b279c Mon Sep 17 00:00:00 2001
From: glados-verma <akashv@google.com>
Date: Wed, 30 Oct 2024 11:03:58 -0700
Subject: [PATCH] Update documentation for the default repeat limit, add
 constant for retrying "forever".

Originally the default was to repeat forever, which was changed to 3 at some point, without updating the documentation.

PiperOrigin-RevId: 691482935
---
 openhtf/core/phase_descriptor.py | 10 ++++--
 openhtf/core/phase_executor.py   |  5 +--
 test/core/phase_executor_test.py | 56 ++++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 4 deletions(-)
 create mode 100644 test/core/phase_executor_test.py

diff --git a/openhtf/core/phase_descriptor.py b/openhtf/core/phase_descriptor.py
index 0dcf0069..7327c1a9 100644
--- a/openhtf/core/phase_descriptor.py
+++ b/openhtf/core/phase_descriptor.py
@@ -24,6 +24,7 @@
 import logging
 import os.path
 import pdb
+import sys
 from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Set, Text, TYPE_CHECKING, Type, Union
 
 import attr
@@ -43,6 +44,10 @@
 if TYPE_CHECKING:
   from openhtf.core import test_state  # pylint: disable=g-import-not-at-top
 
+
+DEFAULT_REPEAT_LIMIT = 3
+MAX_REPEAT_LIMIT = sys.maxsize
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -111,8 +116,9 @@ class PhaseOptions(object):
     repeat_on_measurement_fail: If true, force phase with failed
       measurements to repeat up to repeat_limit times.
     repeat_on_timeout:  If consider repeat on phase timeout, default is No.
-    repeat_limit:  Maximum number of repeats.  None indicates a phase will be
-      repeated infinitely as long as PhaseResult.REPEAT is returned.
+    repeat_limit:  Maximum number of repeats.  DEFAULT_REPEAT_LIMIT applies if
+      this is set to None.  MAX_REPEAT_LIMIT can be used to repeat the phase
+      virtually forever, as long as PhaseResult.REPEAT is returned.
     run_under_pdb: If True, run the phase under the Python Debugger (pdb).  When
       setting this option, increase the phase timeout as well because the
       timeout will still apply when under the debugger.
diff --git a/openhtf/core/phase_executor.py b/openhtf/core/phase_executor.py
index 6c1139bc..f578fda9 100644
--- a/openhtf/core/phase_executor.py
+++ b/openhtf/core/phase_executor.py
@@ -51,7 +51,6 @@
   from openhtf.core import test_state as htf_test_state  # pylint: disable=g-import-not-at-top
 
 DEFAULT_PHASE_TIMEOUT_S = 3 * 60
-DEFAULT_RETRIES = 3
 _JOIN_TRY_INTERVAL_SECONDS = 3
 
 ARG_PARSER = argv.module_parser()
@@ -279,7 +278,9 @@ def execute_phase(
       requested and successfully ran for this phase execution.
     """
     repeat_count = 1
-    repeat_limit = (phase.options.repeat_limit or DEFAULT_RETRIES)
+    repeat_limit = (
+        phase.options.repeat_limit or phase_descriptor.DEFAULT_REPEAT_LIMIT
+    )
     while not self._stopping.is_set():
       is_last_repeat = repeat_count >= repeat_limit
       phase_execution_outcome, profile_stats = self._execute_phase_once(
diff --git a/test/core/phase_executor_test.py b/test/core/phase_executor_test.py
new file mode 100644
index 00000000..86379568
--- /dev/null
+++ b/test/core/phase_executor_test.py
@@ -0,0 +1,56 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for phase_executor module."""
+
+import unittest
+
+import openhtf
+from openhtf.core import phase_descriptor
+
+
+class PhaseExecutorTest(unittest.TestCase):
+
+  def _run_repeating_phase(self, phase_options, expected_call_count):
+    call_count = 0
+
+    @phase_options
+    def repeating_phase():
+      nonlocal call_count
+      call_count += 1
+      if call_count < phase_descriptor.DEFAULT_REPEAT_LIMIT + 1:
+        return phase_descriptor.PhaseResult.REPEAT
+      return phase_descriptor.PhaseResult.STOP
+
+    test = openhtf.Test(repeating_phase)
+    test.execute()
+    self.assertEqual(call_count, expected_call_count)
+
+  def test_execute_phase_with_repeat_limit_unspecified_uses_default_limit(self):
+    self._run_repeating_phase(
+        openhtf.PhaseOptions(),
+        expected_call_count=phase_descriptor.DEFAULT_REPEAT_LIMIT,
+    )
+
+  def test_execute_phase_with_repeat_limit_none_uses_default_limit(self):
+    self._run_repeating_phase(
+        openhtf.PhaseOptions(repeat_limit=None),
+        expected_call_count=phase_descriptor.DEFAULT_REPEAT_LIMIT,
+    )
+
+  def test_execute_phase_with_repeat_limit_max_exceeds_default_limit(self):
+    self._run_repeating_phase(
+        openhtf.PhaseOptions(repeat_limit=phase_descriptor.MAX_REPEAT_LIMIT),
+        expected_call_count=phase_descriptor.DEFAULT_REPEAT_LIMIT + 1,
+    )