From 73cc360c4711abe8e42ea3eab1d537235d4ca340 Mon Sep 17 00:00:00 2001
From: Mia Garrard <mgarrard@meta.com>
Date: Thu, 31 Oct 2024 18:55:01 -0700
Subject: [PATCH] input constructor storage wip

Differential Revision: D65299913
---
 ax/storage/json_store/tests/test_json_store.py | 6 ++++++
 ax/utils/testing/modeling_stubs.py             | 7 +++++++
 2 files changed, 13 insertions(+)

diff --git a/ax/storage/json_store/tests/test_json_store.py b/ax/storage/json_store/tests/test_json_store.py
index b0868eb3d18..24c8a634e57 100644
--- a/ax/storage/json_store/tests/test_json_store.py
+++ b/ax/storage/json_store/tests/test_json_store.py
@@ -208,6 +208,12 @@
             sobol_gpei_generation_node_gs, with_input_constructors_target_trial=True
         ),
     ),
+    (
+        "GenerationStrategy",
+        partial(
+            sobol_gpei_generation_node_gs, with_input_constructors_sq_features=True
+        ),
+    ),
     (
         "GenerationStrategy",
         partial(sobol_gpei_generation_node_gs, with_unlimited_gen_mbm=True),
diff --git a/ax/utils/testing/modeling_stubs.py b/ax/utils/testing/modeling_stubs.py
index bc7d04b33a9..be9943c9663 100644
--- a/ax/utils/testing/modeling_stubs.py
+++ b/ax/utils/testing/modeling_stubs.py
@@ -228,6 +228,7 @@ def sobol_gpei_generation_node_gs(
     with_input_constructors_remaining_n: bool = False,
     with_input_constructors_repeat_n: bool = False,
     with_input_constructors_target_trial: bool = False,
+    with_input_constructors_sq_features: bool = False,
     with_unlimited_gen_mbm: bool = False,
     with_trial_type: bool = False,
     with_is_SOO_transition: bool = False,
@@ -250,6 +251,7 @@ def sobol_gpei_generation_node_gs(
                 with_input_constructors_remaining_n,
                 with_input_constructors_repeat_n,
                 with_input_constructors_target_trial,
+                with_input_constructors_sq_features,
             ]
         )
         > 1
@@ -383,6 +385,11 @@ def sobol_gpei_generation_node_gs(
         sobol_node._input_constructors = {
             purpose: NodeInputConstructors.TARGET_TRIAL_FIXED_FEATURES,
         }
+    elif with_input_constructors_sq_features:
+        purpose = InputConstructorPurpose.STATUS_QUO_FEATURES
+        sobol_node._input_constructors = {
+            purpose: NodeInputConstructors.STATUS_QUO_FEATURES,
+        }
 
     sobol_mbm_GS_nodes = GenerationStrategy(
         name="Sobol+MBM_Nodes",