fix: setup_environment integration test (#1196)

* fix setup_env integration test * standardize logging
basetenlabs · Oct 22, 2024 · ea9aa46 · ea9aa46
1 parent 5906d60
commit ea9aa46
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/truss/templates/server/model_wrapper.py b/truss/templates/server/model_wrapper.py
@@ -469,7 +469,7 @@ async def setup_environment(self, environment: Optional[dict]):
         if not descriptor:
             return
         self._logger.info(
-            f"Executing model.setup_environment with new environment: {environment}"
+            f"Executing model.setup_environment with environment: {environment}"
         )
         if descriptor.is_async:
             return await self._model.setup_environment(environment)

diff --git a/truss/tests/test_model_inference.py b/truss/tests/test_model_inference.py
@@ -792,11 +792,11 @@ def predict(self, model_input):
             ]
         )
         time.sleep(30)
-        single_quote_beta_env_str = beta_env_str.replace('"', "'")
         assert (
-            f"Executing model.setup_environment with environment: {single_quote_beta_env_str}"
+            f"Executing model.setup_environment with environment: {beta_env}"
             in container.logs()
         )
+        single_quote_beta_env_str = beta_env_str.replace('"', "'")
         assert (
             f"setup_environment called with {single_quote_beta_env_str}"
             in container.logs()
@@ -837,11 +837,11 @@ def predict(self, model_input):
             wait_for_server_ready=True,
         )
         # Don't need to wait here because we explicitly grab the environment from dynamic_config_resolver before calling user's load()
-        single_quote_staging_env_str = staging_env_str.replace('"', "'")
         assert (
-            f"Executing model.setup_environment with environment: {single_quote_staging_env_str}"
+            f"Executing model.setup_environment with environment: {staging_env}"
             in container.logs()
         )
+        single_quote_staging_env_str = staging_env_str.replace('"', "'")
         assert (
             f"setup_environment called with {single_quote_staging_env_str}"
             in container.logs()
@@ -860,7 +860,7 @@ def predict(self, model_input):
         )
         time.sleep(30)
         assert (
-            f"Executing model.setup_environment with new environment: {no_env}"
+            f"Executing model.setup_environment with environment: {no_env}"
             in container.logs()
         )
         assert "setup_environment called with None" in container.logs()