Reinforce testenv in ax test case (#2530)

Summary: Pull Request resolved: #2530

Reviewed By: mpolson64

Differential Revision: D58842047

fbshipit-source-id: ebf9d323e7a8ec5bae9e252273fc71b09eabc2fd

ax/storage/sqa_store/utils.py

@@ -38,6 +38,10 @@
     "_steps",
     "analysis_scheduler",
     "_nodes",
+    # ``status_quo_weight_override`` is a field on ``BatchTrial`` not in the
+    # "trial_v2" table
+    # TODO(T193258337)
+    "_status_quo_weight_override",
 }
 SKIP_ATTRS_ERROR_SUFFIX = "Consider adding to COPY_DB_IDS_ATTRS_TO_SKIP if appropriate."
 

ax/utils/common/constants.py

@@ -87,3 +87,6 @@ class Keys(str, Enum):
 
 DEFAULT_WINSORIZATION_LIMITS_MINIMIZATION: Tuple[float, float] = (0.0, 0.2)
 DEFAULT_WINSORIZATION_LIMITS_MAXIMIZATION: Tuple[float, float] = (0.2, 0.0)
+
+TESTENV_ENV_KEY = "TESTENV"
+TESTENV_ENV_VAL = "True"

ax/utils/common/testutils.py

@@ -14,6 +14,7 @@
 import io
 import linecache
 import logging
+import os
 import signal
 import sys
 import types
@@ -39,6 +40,7 @@
 import numpy as np
 from ax.exceptions.core import AxParameterWarning
 from ax.utils.common.base import Base
+from ax.utils.common.constants import TESTENV_ENV_KEY, TESTENV_ENV_VAL
 from ax.utils.common.equality import object_attribute_dicts_find_unequal_fields
 from ax.utils.common.logger import get_logger
 from botorch.exceptions.warnings import InputDataWarning
@@ -301,6 +303,12 @@ def signal_handler(signum: int, frame: Optional[FrameType]) -> None:
 
         super().__init__(methodName=methodName)
         signal.signal(signal.SIGALRM, signal_handler)
+        # This is set to indicate we are running in a test environment.  Code can check
+        # this to:
+        # * more strictly enforce SQL encoding
+        #  (https://github.com/facebook/Ax/blob/main/ax/storage/sqa_store/save.py#L598)
+        # * avoid actions that will affect product environments
+        os.environ[TESTENV_ENV_KEY] = TESTENV_ENV_VAL
 
     def setUp(self) -> None:
         """