Address PR comments

Author: Chia-Eng

src/sagemaker/estimator.py

@@ -1808,6 +1808,8 @@ def __init__(
                 ``disable_profiler`` parameter to ``True``.
             disable_profiler (bool): Specifies whether Debugger monitoring and profiling
                 will be disabled (default: ``False``).
+            environment (dict[str, str]) : A string to string map contains environment
+                variables to set in the Docker container.
         """
         self.image_uri = image_uri
         self.hyperparam_dict = hyperparameters.copy() if hyperparameters else {}

tests/unit/test_estimator.py

@@ -2697,24 +2697,7 @@ def test_add_environment_variables_to_train_args(sagemaker_session):
 
     sagemaker_session.train.assert_called_once()
     args = sagemaker_session.train.call_args[1]
-    assert args["environment"]
-
-def test_no_environment_variables_in_train_args(sagemaker_session):
-    e = Estimator(
-        IMAGE_URI,
-        ROLE,
-        INSTANCE_COUNT,
-        INSTANCE_TYPE,
-        output_path=OUTPUT_PATH,
-        sagemaker_session=sagemaker_session,
-        environment=None,
-    )
-
-    e.fit()
-
-    sagemaker_session.train.assert_called_once()
-    args = sagemaker_session.train.call_args[1]
-    assert not args["environment"]
+    assert args["environment"] == ENV_INPUT
 
 def test_generic_to_fit_with_sagemaker_metrics_enabled(sagemaker_session):
     e = Estimator(