code: def test_ModelBridge(
      self,
      mock_fit: Mock,
      mock_gen_arms: Mock,
      mock_observations_from_data: Mock,
  ) -> None:
      # Test that on init transforms are stored and applied in the correct order
      transforms = [transform_1, transform_2]
      exp = get_experiment_for_value()
      ss = get_search_space_for_value()
      modelbridge = ModelBridge(
          search_space=ss,
          model=Model(),
          # pyre-fixme[6]: For 3rd param expected
          #  `Optional[List[Type[Transform]]]` but got `List[Type[Union[transform_1,
          #  transform_2]]]`.
          transforms=transforms,
          experiment=exp,
          # pyre-fixme[6]: For 5th param expected `Optional[Data]` but got `int`.
          data=0,
      )
      self.assertFalse(
          modelbridge._experiment_has_immutable_search_space_and_opt_config
      )
      self.assertEqual(
          list(modelbridge.transforms.keys()),
          ["Cast", "transform_1", "transform_2"],
      )
      fit_args = mock_fit.mock_calls[0][2]
      self.assertTrue(
          fit_args["search_space"] == get_search_space_for_value(8.0)
      )
      self.assertTrue(fit_args["observations"] == [])
      self.assertTrue(mock_observations_from_data.called)
  
      # Test deprecation error on update.
      with self.assertRaisesRegex(
          DeprecationWarning, "ModelBridge.update"
      ):
          modelbridge.update(Mock(), Mock())
  
      # Test prediction on out of design features.
      modelbridge._predict = mock.MagicMock(
          "ax.modelbridge.base.ModelBridge._predict",
          autospec=True,
          side_effect=ValueError("Out of Design"),
      )
      # This point is in design, and thus failures in predict are legitimate.
      with mock.patch.object(
          ModelBridge,
          "model_space",
          return_value=get_search_space_for_range_values,
      ):
          with self.assertRaises(ValueError):
              modelbridge.predict([get_observation2().features])
  
      # This point is out of design, and not in training data.
      with self.assertRaises(ValueError):
          modelbridge.predict([get_observation_status_quo0().features])
  
      # Now it's in the training data.
      with mock.patch.object(
          ModelBridge,
          "get_training_data",
          return_value=[get_observation_status_quo0()],
      ):
          # Return raw training value.
          self.assertEqual(
              modelbridge.predict(
                  [get_observation_status_quo0().features]
              ),
              unwrap_observation_data(
                  [get_observation_status_quo0().data]
              ),
          )
  
      # Test that transforms are applied correctly on predict
      modelbridge._predict = mock.MagicMock(
          "ax.modelbridge.base.ModelBridge._predict",
          autospec=True,
          return_value=[get_observation2trans().data],
      )
      modelbridge.predict([get_observation2().features])
      # Observation features sent to _predict are un-transformed afterwards
      # pyre-fixme[16]: Callable `_predict` has no attribute `assert_called_with`.
      modelbridge._predict.assert_called_with(
          [get_observation2().features]
      )
  
      # Check that _single_predict is equivalent here.
      modelbridge._single_predict([get_observation2().features])
      # Observation features sent to _predict are un-transformed afterwards
      modelbridge._predict.assert_called_with(
          [get_observation2().features]
      )
  
      # Test transforms applied on gen
      modelbridge._gen = mock.MagicMock(
          "ax.modelbridge.base.ModelBridge._gen",
          autospec=True,
          return_value=GenResults(
              observation_features=[get_observation1trans().features],
              weights=[2],
          ),
      )
      oc = OptimizationConfig(
          objective=Objective(metric=Metric(name="test_metric"))
      )
      modelbridge._set_kwargs_to_save(
          model_key="TestModel", model_kwargs={}, bridge_kwargs={}
      )
      # Test input error when generating 0 candidates.
      with self.assertRaisesRegex(
          UserInputError, "Attempted to generate"
      ):
          modelbridge.gen(n=0)
      gr = modelbridge.gen(
          n=1,
          search_space=get_search_space_for_value(),
          optimization_config=oc,
          pending_observations={"a": [get_observation2().features]},
          fixed_features=ObservationFeatures({"x": 5}),
      )
      self.assertEqual(gr._model_key, "TestModel")
      # pyre-fixme[16]: Callable `_gen` has no attribute `assert_called_with`.
      modelbridge._gen.assert_called_with(
          n=1,
          search_space=SearchSpace(
              [FixedParameter("x", ParameterType.FLOAT, 8.0)]
          ),
          optimization_config=oc,
          pending_observations={
              "a": [get_observation2trans().features]
          },
          fixed_features=ObservationFeatures({"x": 36}),
          model_gen_options=None,
      )
      mock_gen_arms.assert_called_with(
          arms_by_signature={},
          observation_features=[get_observation1().features],
      )
  
      # Gen with no pending observations and no fixed features
      modelbridge.gen(
          n=1,
          search_space=get_search_space_for_value(),
          optimization_config=None,
      )
      modelbridge._gen.assert_called_with(
          n=1,
          search_space=SearchSpace(
              [FixedParameter("x", ParameterType.FLOAT, 8.0)]
          ),
          optimization_config=None,
          pending_observations={},
          fixed_features=None,
          model_gen_options=None,
      )
  
      # Gen with multi-objective optimization config.
      oc2 = OptimizationConfig(
          objective=ScalarizedObjective(
              metrics=[
                  Metric(name="test_metric"),
                  Metric(name="test_metric_2"),
              ]
          )
      )
      modelbridge.gen(
          n=1,
          search_space=get_search_space_for_value(),
          optimization_config=oc2,
      )
      modelbridge._gen.assert_called_with(
          n=1,
          search_space=SearchSpace(
              [FixedParameter("x", ParameterType.FLOAT, 8.0)]
          ),
          optimization_config=oc2,
          pending_observations={},
          fixed_features=None,
          model_gen_options=None,
      )
  
      # Test transforms applied on cross_validate and the warning is suppressed.
      called = False
  
      def warn_and_return_mock_obs(
          *args: Any, **kwargs: Any
      ) -> List[ObservationData]:
          nonlocal called
          called = True
          warnings.warn(
              "Input data is not standardized. Please consider scaling the "
              "input to zero mean and unit variance.",
              InputDataWarning,
          )
          return [get_observation1trans().data]
  
      modelbridge._cross_validate = mock.MagicMock(
          "ax.modelbridge.base.ModelBridge._cross_validate",
          autospec=True,
          side_effect=warn_and_return_mock_obs,
      )
      cv_training_data = [get_observation2()]
      cv_test_points = [get_observation1().features]
      with warnings.catch_warnings(record=True) as ws:
          cv_predictions = modelbridge.cross_validate(
              cv_training_data=cv_training_data,
              cv_test_points=cv_test_points,
          )
      self.assertTrue(called)
      self.assertFalse(any(w.category is InputDataWarning for w in ws))
      # pyre-fixme[16]: Callable `_cross_validate` has no attribute
      #  `assert_called_with`.
      modelbridge._cross_validate.assert_called_with(
          search_space=SearchSpace(
              [FixedParameter("x", ParameterType.FLOAT, 8.0)]
          ),
          cv_training_data=[get_observation2trans()],
          cv_test_points=[
              get_observation1().features
          ],  # untransformed after
      )
      self.assertTrue(cv_predictions == [get_observation1().data])
  
      # Test stored training data
      obs = modelbridge.get_training_data()
      self.assertTrue(obs == [get_observation1(), get_observation2()])
      self.assertEqual(modelbridge.metric_names, {"a", "b"})
      self.assertIsNone(modelbridge.status_quo)
      self.assertTrue(
          modelbridge.model_space == get_search_space_for_value()
      )
      self.assertEqual(modelbridge.training_in_design, [False, False])
  
      with self.assertRaises(ValueError):
          modelbridge.training_in_design = [True, True, False]
  
      with self.assertRaises(ValueError):
          modelbridge.training_in_design = [True, True, False]
  
      # Test feature_importances
      with self.assertRaises(NotImplementedError):
          modelbridge.feature_importances("a")
  
      # Test transform observation features
      with mock.patch(
          "ax.modelbridge.base.ModelBridge._transform_observation_features",
          autospec=True,
      ) as mock_tr:
          modelbridge.transform_observation_features(
              [get_observation2().features]
          )
      mock_tr.assert_called_with(
          modelbridge, [get_observation2trans().features]
      )
  
      # Test that fit is not called when fit_on_init = False.
      mock_fit.reset_mock()
      modelbridge = ModelBridge(
          search_space=ss,
          model=Model(),
          fit_on_init=False,
      )
      self.assertEqual(mock_fit.call_count, 0)
  
      # Test error when fit_tracking_metrics is False and optimization
      # config is not specified.
      with self.assertRaisesRegex(
          UserInputError, "fit_tracking_metrics"
      ):
          ModelBridge(
              search_space=ss,
              model=Model(),
              fit_tracking_metrics=False,
          )
  
      # Test error when fit_tracking_metrics is False and optimization
      # config is updated to include new metrics.
      modelbridge = ModelBridge(
          search_space=ss,
          model=Model(),
          optimization_config=oc,
          fit_tracking_metrics=False,
      )
      new_oc = OptimizationConfig(
          objective=Objective(metric=Metric(name="test_metric2"))
      )
      with self.assertRaisesRegex(
          UnsupportedError, "fit_tracking_metrics"
      ):
          modelbridge.gen(n=1, optimization_config=new_oc)
  
imports: import warnings
  from typing import Any, List
  from unittest import mock
  from unittest.mock import Mock
  import numpy
  from ax.core.arm import Arm
  from ax.core.data import Data
  from ax.core.metric import Metric
  from ax.core.objective import Objective, ScalarizedObjective
  from ax.core.observation import ObservationData, ObservationFeatures
  from ax.core.optimization_config import OptimizationConfig
  from ax.core.parameter import FixedParameter, ParameterType, RangeParameter
  from ax.core.search_space import SearchSpace
  from ax.exceptions.core import UnsupportedError, UserInputError
  from ax.modelbridge.base import clamp_observation_features, gen_arms, GenResults, ModelBridge, unwrap_observation_data
  from ax.modelbridge.registry import Models, Y_trans
  from ax.models.base import Model
  from ax.utils.common.constants import Keys
  from ax.utils.common.testutils import TestCase
  from ax.utils.testing.core_stubs import get_branin_experiment, get_branin_experiment_with_multi_objective, get_experiment, get_experiment_with_repeated_arms, get_non_monolithic_branin_moo_data, get_optimization_config_no_constraints, get_search_space_for_range_value, get_search_space_for_range_values, get_search_space_for_value
  from ax.utils.testing.modeling_stubs import get_experiment_for_value, get_observation1, get_observation1trans, get_observation2, get_observation2trans, get_observation_status_quo0, get_observation_status_quo1, transform_1, transform_2
  from botorch.exceptions.warnings import InputDataWarning
label_desc: Use image node_id 1 for calling the BaseModelBridgeTest obj's underlying member method code with example usage: obj.test_ModelBridge(mock_fit, mock_gen_arms, mock_observations_from_data) and returns: list