updated as per review comments

Author: VipulMascarenhas

ads/aqua/__init__.py

@@ -8,7 +8,7 @@
 
 from ads import logger, set_auth
 from ads.aqua.utils import fetch_service_compartment
-from ads.config import NB_SESSION_OCID, OCI_RESOURCE_PRINCIPAL_VERSION
+from ads.config import OCI_RESOURCE_PRINCIPAL_VERSION
 
 ENV_VAR_LOG_LEVEL = "ADS_AQUA_LOG_LEVEL"
 
@@ -36,8 +36,3 @@ def set_log_level(log_level: str):
 ODSC_MODEL_COMPARTMENT_OCID = (
     os.environ.get("ODSC_MODEL_COMPARTMENT_OCID") or fetch_service_compartment()
 )
-if not ODSC_MODEL_COMPARTMENT_OCID:
-    if NB_SESSION_OCID:
-        logger.error(
-            f"Aqua is not available for this notebook session {NB_SESSION_OCID}."
-        )

ads/aqua/cli.py

@@ -6,11 +6,17 @@
 import os
 import sys
 
-from ads.aqua import ENV_VAR_LOG_LEVEL, set_log_level, ODSC_MODEL_COMPARTMENT_OCID
+from ads.aqua import (
+    ENV_VAR_LOG_LEVEL,
+    set_log_level,
+    ODSC_MODEL_COMPARTMENT_OCID,
+    logger,
+)
 from ads.aqua.deployment import AquaDeploymentApp
 from ads.aqua.evaluation import AquaEvaluationApp
 from ads.aqua.finetune import AquaFineTuningApp
 from ads.aqua.model import AquaModelApp
+from ads.config import NB_SESSION_OCID
 
 
 class AquaCommand:
@@ -44,4 +50,11 @@ def __init__(
         set_log_level(log_level)
         # gracefully exit if env var is not set
         if not ODSC_MODEL_COMPARTMENT_OCID:
-            sys.exit(0)
+            logger.error(
+                "ODSC_MODEL_COMPARTMENT_OCID environment variable is not set for Aqua."
+            )
+            if NB_SESSION_OCID:
+                logger.error(
+                    f"Aqua is not available for the notebook session {NB_SESSION_OCID}."
+                )
+            sys.exit(1)

ads/aqua/utils.py

@@ -582,7 +582,7 @@ def fetch_service_compartment() -> Union[str, None]:
         )
     except AquaFileNotFoundError:
         logger.error(
-            f"ODSC_MODEL_COMPARTMENT_OCID environment variable is not set for Aqua."
+            f"Config file {config_file_name}/{CONTAINER_INDEX} to fetch service compartment OCID could not be found."
         )
         return
     compartment_mapping = config.get(COMPARTMENT_MAPPING_KEY)

tests/unitary/with_extras/aqua/test_cli.py

@@ -7,7 +7,6 @@
 import os
 import logging
 import subprocess
-import pytest
 from unittest import TestCase
 from unittest.mock import patch
 from importlib import reload
@@ -21,7 +20,7 @@
 class TestAquaCLI(TestCase):
     """Tests the AQUA CLI."""
 
-    DEFAUL_AQUA_CLI_LOGGING_LEVEL = "ERROR"
+    DEFAULT_AQUA_CLI_LOGGING_LEVEL = "ERROR"
     logger = logging.getLogger(__name__)
     logging.basicConfig(
         format="%(asctime)s %(module)s %(levelname)s: %(message)s",
@@ -30,11 +29,6 @@ class TestAquaCLI(TestCase):
     )
     SERVICE_COMPARTMENT_ID = "ocid1.compartment.oc1..<OCID>"
 
-    def setUp(self):
-        os.environ["ODSC_MODEL_COMPARTMENT_OCID"] = TestAquaCLI.SERVICE_COMPARTMENT_ID
-        reload(ads.aqua)
-        reload(ads.aqua.cli)
-
     def test_entrypoint(self):
         """Tests CLI entrypoint."""
         result = subprocess.run(["ads", "aqua", "--help"], capture_output=True)
@@ -43,23 +37,55 @@ def test_entrypoint(self):
 
     @parameterized.expand(
         [
-            ("default", None, DEFAUL_AQUA_CLI_LOGGING_LEVEL),
+            ("default", None, DEFAULT_AQUA_CLI_LOGGING_LEVEL),
             ("set logging level", "info", "info"),
         ]
     )
-    @patch("ads.aqua.cli.set_log_level")
-    def test_aquacommand(self, name, arg, expected, mock_setting_log):
-        """Tests aqua command initailzation."""
-        if arg:
-            AquaCommand(arg)
-        else:
-            AquaCommand()
-        mock_setting_log.assert_called_with(expected)
+    def test_aquacommand(self, name, arg, expected):
+        """Tests aqua command initialization."""
+        with patch.dict(
+            os.environ,
+            {"ODSC_MODEL_COMPARTMENT_OCID": TestAquaCLI.SERVICE_COMPARTMENT_ID},
+        ):
+            reload(ads.config)
+            reload(ads.aqua)
+            reload(ads.aqua.cli)
+            with patch("ads.aqua.cli.set_log_level") as mock_setting_log:
+                if arg:
+                    AquaCommand(arg)
+                else:
+                    AquaCommand()
+                mock_setting_log.assert_called_with(expected)
+
+    @parameterized.expand(
+        [
+            ("default", None),
+            ("using jupyter instance", "nb-session-ocid"),
+        ]
+    )
+    def test_aqua_command_without_compartment_env_var(self, name, session_ocid):
+        """Test whether exit is called when ODSC_MODEL_COMPARTMENT_OCID is not set. Also check if NB_SESSION_OCID is
+        set then log the appropriate message."""
 
-    @patch("sys.exit")
-    def test_aqua_command_without_compartment_env_var(self, mock_exit):
-        os.environ.pop("ODSC_MODEL_COMPARTMENT_OCID", None)
-        reload(ads.aqua)
-        reload(ads.aqua.cli)
-        AquaCommand()
-        mock_exit.assert_called_with(0)
+        with patch("sys.exit") as mock_exit:
+            env_dict = {"ODSC_MODEL_COMPARTMENT_OCID": ""}
+            if session_ocid:
+                env_dict.update({"NB_SESSION_OCID": session_ocid})
+            with patch.dict(os.environ, env_dict):
+                reload(ads.config)
+                reload(ads.aqua)
+                reload(ads.aqua.cli)
+                with patch("ads.aqua.cli.set_log_level") as mock_setting_log:
+                    with patch("ads.aqua.logger.error") as mock_logger_error:
+                        AquaCommand()
+                        mock_setting_log.assert_called_with(
+                            TestAquaCLI.DEFAULT_AQUA_CLI_LOGGING_LEVEL
+                        )
+                        mock_logger_error.assert_any_call(
+                            "ODSC_MODEL_COMPARTMENT_OCID environment variable is not set for Aqua."
+                        )
+                        if session_ocid:
+                            mock_logger_error.assert_any_call(
+                                f"Aqua is not available for the notebook session {session_ocid}."
+                            )
+                        mock_exit.assert_called_with(1)