From ff1445aa089f47afd86079057689970d19d85cd6 Mon Sep 17 00:00:00 2001
From: Daniel Cohen <danielcohennyc@meta.com>
Date: Thu, 27 Jun 2024 14:11:02 -0700
Subject: [PATCH] Reinforce testenv in ax test case (#2530)

Summary: Pull Request resolved: https://github.com/facebook/Ax/pull/2530

Reviewed By: mpolson64

Differential Revision: D58842047

fbshipit-source-id: ebf9d323e7a8ec5bae9e252273fc71b09eabc2fd
---
 ax/storage/sqa_store/utils.py | 4 ++++
 ax/utils/common/constants.py  | 3 +++
 ax/utils/common/testutils.py  | 8 ++++++++
 3 files changed, 15 insertions(+)

diff --git a/ax/storage/sqa_store/utils.py b/ax/storage/sqa_store/utils.py
index 76ab381c247..8aedada8899 100644
--- a/ax/storage/sqa_store/utils.py
+++ b/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."
 
diff --git a/ax/utils/common/constants.py b/ax/utils/common/constants.py
index d3eb9107a2e..5cc3d34a43c 100644
--- a/ax/utils/common/constants.py
+++ b/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"
diff --git a/ax/utils/common/testutils.py b/ax/utils/common/testutils.py
index 930733c3559..a9d8239e55c 100644
--- a/ax/utils/common/testutils.py
+++ b/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:
         """