Draft fix

comet-ml · Sep 16, 2024 · 2dd2b31 · 2dd2b31
1 parent 7f0ac9b
commit 2dd2b31
Show file tree

Hide file tree

Showing 3 changed files with 169 additions and 69 deletions.
diff --git a/sdks/python/src/opik/config.py b/sdks/python/src/opik/config.py
@@ -9,7 +9,12 @@
 
 from . import dict_utils
 
-PathType = Union[pathlib.Path, str, List[Union[pathlib.Path, str]], Tuple[Union[pathlib.Path, str], ...]]
+PathType = Union[
+    pathlib.Path,
+    str,
+    List[Union[pathlib.Path, str]],
+    Tuple[Union[pathlib.Path, str], ...],
+]
 
 _SESSION_CACHE_DICT: Dict[str, Any] = {}
 
@@ -19,7 +24,7 @@
 OPIK_PROJECT_DEFAULT_NAME: Final[str] = "Default Project"
 OPIK_WORKSPACE_DEFAULT_NAME: Final[str] = "default"
 
-CONFIG_FILE_PATH_DEFAULT: Final[str] = '~/.opik.config'
+CONFIG_FILE_PATH_DEFAULT: Final[str] = "~/.opik.config"
 
 LOGGER = logging.getLogger(__name__)
 
@@ -39,7 +44,9 @@ def __init__(
     def _read_file(self, file_path: pathlib.Path) -> dict[str, Any]:
         config = configparser.ConfigParser()
         config.read(file_path)
-        config_values = {section: dict(config.items(section)) for section in config.sections()}
+        config_values = {
+            section: dict(config.items(section)) for section in config.sections()
+        }
 
         if "opik" in config_values:
             return config_values["opik"]
@@ -151,13 +158,14 @@ def save_to_file(self) -> None:
         if self.api_key is not None:
             config_file_content["opik"]["api_key"] = self.api_key
 
-        with open(self.config_file_fullpath, mode="w+", encoding="utf-8") as config_file:
+        with open(
+            self.config_file_fullpath, mode="w+", encoding="utf-8"
+        ) as config_file:
             config_file_content.write(config_file)
 
         LOGGER.info(f"Saved configuration to a file: {self.config_file_fullpath}")
 
 
-
 def update_session_config(key: str, value: Any) -> None:
     _SESSION_CACHE_DICT[key] = value
 

diff --git a/sdks/python/src/opik/opik_login.py b/sdks/python/src/opik/opik_login.py
@@ -1,30 +1,31 @@
-"""
-
-"""
+""" """
 
 import logging
-import sys
 from getpass import getpass
 from typing import Final, Optional
 
 import httpx
 
 import opik.config
 from opik import httpx_client
-from opik.config import OPIK_BASE_URL_LOCAL, OPIK_BASE_URL_CLOUD, OPIK_WORKSPACE_DEFAULT_NAME
+from opik.config import (
+    OPIK_BASE_URL_LOCAL,
+    OPIK_BASE_URL_CLOUD,
+    OPIK_WORKSPACE_DEFAULT_NAME,
+)
 from opik.exceptions import ConfigurationError
 
 LOGGER = logging.getLogger(__name__)
 
-HEALTH_CHECK_URL_POSTFIX: Final[str] = '/is-alive/ping'
+HEALTH_CHECK_URL_POSTFIX: Final[str] = "/is-alive/ping"
 HEALTH_CHECK_TIMEOUT: Final[float] = 1.0
 
 
 def is_interactive() -> bool:
     """
     Returns True if in interactive mode
     """
-    #return bool(getattr(sys, "ps1", sys.flags.interactive))
+    # return bool(getattr(sys, "ps1", sys.flags.interactive))
     return True
 
 
@@ -58,9 +59,11 @@ def is_workspace_name_correct(api_key: str, workspace: str) -> bool:
     url = "https://www.comet.com/api/rest/v2/workspaces"
 
     client = httpx.Client()
-    client.headers.update({
-        "Authorization": f"{api_key}",
-    })
+    client.headers.update(
+        {
+            "Authorization": f"{api_key}",
+        }
+    )
 
     try:
         response = client.get(url=url)
@@ -88,9 +91,11 @@ def is_api_key_correct(api_key: str) -> bool:
     url = "https://www.comet.com/api/rest/v2/account-details"
 
     client = httpx.Client()
-    client.headers.update({
-        "Authorization": f"{api_key}",
-    })
+    client.headers.update(
+        {
+            "Authorization": f"{api_key}",
+        }
+    )
 
     try:
         response = client.get(url=url)
@@ -116,25 +121,29 @@ def get_default_workspace(api_key: str) -> str:
     url = "https://www.comet.com/api/rest/v2/account-details"
 
     client = httpx.Client()
-    client.headers.update({
-        "Authorization": f"{api_key}",
-    })
+    client.headers.update(
+        {
+            "Authorization": f"{api_key}",
+        }
+    )
 
     try:
         response = client.get(url=url)
     except Exception as e:
         raise ConnectionError(f"Error while getting default workspace name: {str(e)}")
 
     if response.status_code != 200:
-        raise ConnectionError(f"Error while getting default workspace name: {response.text}")
+        raise ConnectionError(
+            f"Error while getting default workspace name: {response.text}"
+        )
 
     return response.json()["defaultWorkspaceName"]
 
 
 def _update_config(
-        api_key: Optional[str],
-        url: str,
-        workspace: str,
+    api_key: Optional[str],
+    url: str,
+    workspace: str,
 ) -> None:
     """
     Save changes to config file and update current session config
@@ -183,7 +192,9 @@ def _ask_for_url() -> str:
             LOGGER.error("Wrong URL. Please try again.")
             retries -= 1
 
-    raise ConfigurationError("Can't use URL provided by user. Opik instance is not active or not found.")
+    raise ConfigurationError(
+        "Can't use URL provided by user. Opik instance is not active or not found."
+    )
 
 
 def _ask_for_api_key() -> str:
@@ -211,7 +222,9 @@ def _ask_for_workspace(api_key: str) -> str:
     retries = 3
 
     while retries > 0:
-        user_input_workspace = input("Please enter your cloud Opik instance workspace name: ")
+        user_input_workspace = input(
+            "Please enter your cloud Opik instance workspace name: "
+        )
 
         if is_workspace_name_correct(api_key, user_input_workspace):
             return user_input_workspace
@@ -237,10 +250,10 @@ def ask_user_for_approval(message: str) -> bool:
 
 
 def login(
-        api_key: Optional[str] = None,
-        url: Optional[str] = None,
-        workspace: Optional[str] = None,
-        use_local: bool = False
+    api_key: Optional[str] = None,
+    url: Optional[str] = None,
+    workspace: Optional[str] = None,
+    use_local: bool = False,
 ) -> None:
     """
     Args:
@@ -267,8 +280,8 @@ def login(
 
 
 def _login_cloud(
-        api_key: Optional[str],
-        workspace: Optional[str],
+    api_key: Optional[str],
+    workspace: Optional[str],
 ) -> None:
     """
     Login to cloud Opik instance
@@ -280,35 +293,51 @@ def _login_cloud(
     if is_interactive() is False and api_key is None and current_config.api_key is None:
         raise ConfigurationError("No API key provided for cloud Opik instance.")
 
-    if is_interactive() is False and workspace is None and current_config.workspace is None:
+    if (
+        is_interactive() is False
+        and workspace is None
+        and current_config.workspace is None
+    ):
         raise ConfigurationError("No workspace name provided for cloud Opik instance.")
 
     # Ask for API key
     if api_key is None and current_config.api_key is None:
-        LOGGER.info("You can find your API key here: https://www.comet.com/api/my/settings/")
+        LOGGER.info(
+            "You can find your API key here: https://www.comet.com/api/my/settings/"
+        )
         api_key = _ask_for_api_key()
         config_file_needs_updating = True
     elif api_key is None and current_config.api_key is not None:
         api_key = current_config.api_key
 
-    # if workspace already configured - will use this value
-    if "workspace" in current_config.model_fields_set and current_config.workspace != OPIK_WORKSPACE_DEFAULT_NAME:
-        workspace = current_config.workspace
-
-    # Check what their default workspace is, and we ask them if they want to use the default workspace
-    if workspace is None:
-        default_workspace = get_default_workspace(api_key)
-        use_default_workspace = ask_user_for_approval(f"Do you want to use \"{default_workspace}\" workspace? Y/n: ")
 
-        if use_default_workspace:
+    # Check passed workspace (if it was passed)
+    if workspace is not None:
+        if is_workspace_name_correct(api_key, workspace):
             config_file_needs_updating = True
-            workspace = default_workspace
+        else:
+            raise ConfigurationError("Workspace `%s` is incorrect for the given API key.", workspace)
+    else:
+        # Workspace was not passed, we check if there is already configured value
+        # if workspace already configured - will use this value
+        if (
+            "workspace" in current_config.model_fields_set
+            and current_config.workspace != OPIK_WORKSPACE_DEFAULT_NAME
+        ):
+            workspace = current_config.workspace
+
+        # Check what their default workspace is, and we ask them if they want to use the default workspace
+        if workspace is None:
+            default_workspace = get_default_workspace(api_key)
+            use_default_workspace = ask_user_for_approval(
+                f'Do you want to use "{default_workspace}" workspace? Y/n: '
+            )
 
-            if not is_workspace_name_correct(api_key, workspace):
-                LOGGER.warning("Workspace name is incorrect.")
+            if use_default_workspace:
+                workspace = default_workspace
+            else:
                 workspace = _ask_for_workspace(api_key=api_key)
-        else:
-            workspace = _ask_for_workspace(api_key=api_key)
+                config_file_needs_updating = True                
 
     if config_file_needs_updating:
         _update_config(
@@ -351,7 +380,8 @@ def _login_local(url: Optional[str]) -> None:
             return
 
         use_url = ask_user_for_approval(
-            f"Found local Opik instance on: {OPIK_BASE_URL_LOCAL}\nDo you want to use it? Y/n: ")
+            f"Found local Opik instance on: {OPIK_BASE_URL_LOCAL}\nDo you want to use it? Y/n: "
+        )
 
         if use_url:
             _update_config(

diff --git a/sdks/python/tests/unit/test_opik_login.py b/sdks/python/tests/unit/test_opik_login.py
@@ -7,28 +7,90 @@
 @pytest.mark.parametrize(
     "api_key, url, workspace, local, should_raise",
     [
-        (None, "http://example.com", "workspace1", True, False),  # Missing api_key, local=True
-        (None, "http://example.com", "workspace1", False, True),  # Missing api_key, local=False
-        ("apikey123", None, "workspace1", True, True),            # Missing url, local=True
-        ("apikey123", None, "workspace1", False, True),           # Missing url, local=False
-        ("apikey123", "http://example.com", None, True, True),    # Missing workspace, local=True
-        ("apikey123", "http://example.com", None, False, True),   # Missing workspace, local=False
-        (None, None, "workspace1", True, True),                   # Missing api_key and url, local=True
-        (None, None, "workspace1", False, True),                  # Missing api_key and url, local=False
-        (None, "http://example.com", None, True, True),           # Missing api_key and workspace, local=True
-        (None, "http://example.com", None, False, True),          # Missing api_key and workspace, local=False
-        ("apikey123", None, None, True, True),                    # Missing url and workspace, local=True
-        ("apikey123", None, None, False, True),                   # Missing url and workspace, local=False
-        (None, None, None, True, True),                           # All missing, local=True
-        (None, None, None, False, True),                          # All missing, local=False
-        ("apikey123", "http://example.com", "workspace1", True, False),  # All present, local=True
-        ("apikey123", "http://example.com", "workspace1", False, False), # All present, local=False
-    ]
+        (
+            None,
+            "http://example.com",
+            "workspace1",
+            True,
+            False,
+        ),  # Missing api_key, local=True
+        (
+            None,
+            "http://example.com",
+            "workspace1",
+            False,
+            True,
+        ),  # Missing api_key, local=False
+        ("apikey123", None, "workspace1", True, True),  # Missing url, local=True
+        ("apikey123", None, "workspace1", False, True),  # Missing url, local=False
+        (
+            "apikey123",
+            "http://example.com",
+            None,
+            True,
+            True,
+        ),  # Missing workspace, local=True
+        (
+            "apikey123",
+            "http://example.com",
+            None,
+            False,
+            True,
+        ),  # Missing workspace, local=False
+        (None, None, "workspace1", True, True),  # Missing api_key and url, local=True
+        (None, None, "workspace1", False, True),  # Missing api_key and url, local=False
+        (
+            None,
+            "http://example.com",
+            None,
+            True,
+            True,
+        ),  # Missing api_key and workspace, local=True
+        (
+            None,
+            "http://example.com",
+            None,
+            False,
+            True,
+        ),  # Missing api_key and workspace, local=False
+        ("apikey123", None, None, True, True),  # Missing url and workspace, local=True
+        (
+            "apikey123",
+            None,
+            None,
+            False,
+            True,
+        ),  # Missing url and workspace, local=False
+        (None, None, None, True, True),  # All missing, local=True
+        (None, None, None, False, True),  # All missing, local=False
+        (
+            "apikey123",
+            "http://example.com",
+            "workspace1",
+            True,
+            False,
+        ),  # All present, local=True
+        (
+            "apikey123",
+            "http://example.com",
+            "workspace1",
+            False,
+            False,
+        ),  # All present, local=False
+    ],
 )
 def test_login__force_new_settings__fail(api_key, url, workspace, local, should_raise):
     if should_raise:
         with pytest.raises(ConfigurationError):
-            login(api_key=api_key, url=url, workspace=workspace, force=True, use_local=local)
+            login(
+                api_key=api_key,
+                url=url,
+                workspace=workspace,
+                force=True,
+                use_local=local,
+            )
     else:
         # No exception should be raised
-        login(api_key=api_key, url=url, workspace=workspace, force=True, use_local=local)
+        login(
+            api_key=api_key, url=url, workspace=workspace, force=True, use_local=local
+        )