Add support for studio_token and studio_repo_url to be provided a…

…s args. (#29)
iterative · Apr 20, 2023 · 66e25cf · 66e25cf
1 parent 0b86738
commit 66e25cf
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 5 deletions.
diff --git a/src/dvc_studio_client/post_live_metrics.py b/src/dvc_studio_client/post_live_metrics.py
@@ -53,15 +53,15 @@ def get_studio_repo_url() -> Optional[str]:
         return None
 
 
-def get_studio_token_and_repo_url():
-    studio_token = getenv(DVC_STUDIO_TOKEN) or getenv(STUDIO_TOKEN)
+def get_studio_token_and_repo_url(studio_token=None, studio_repo_url=None):
+    studio_token = studio_token or getenv(DVC_STUDIO_TOKEN) or getenv(STUDIO_TOKEN)
     if studio_token is None:
         logger.debug(
             f"{DVC_STUDIO_TOKEN} not found. Skipping `post_studio_live_metrics`"
         )
         return None, None
 
-    studio_repo_url = getenv(STUDIO_REPO_URL, None)
+    studio_repo_url = studio_repo_url or getenv(STUDIO_REPO_URL, None)
     if studio_repo_url is None:
         logger.debug(f"`{STUDIO_REPO_URL}` not found. Trying to automatically find it.")
         studio_repo_url = get_studio_repo_url()
@@ -79,6 +79,8 @@ def post_live_metrics(
     params: Optional[Dict[str, Any]] = None,
     plots: Optional[Dict[str, Any]] = None,
     step: Optional[int] = None,
+    studio_token: Optional[str] = None,
+    studio_repo_url: Optional[str] = None,
 ) -> Optional[bool]:
     """Post `event_type` to Studio's `api/live`.
 
@@ -143,14 +145,18 @@ def post_live_metrics(
             Usually comes from DVCLive `Live.step` property.
             Required in when `event_type="data"`.
             Defaults to `None`.
-
+        studio_token (Optional[str]): Studio access token obtained from the UI.
+        studio_repo_url (Optional[str]): URL of the Git repository that has been
+            imported into Studio UI.
     Returns:
         Optional[bool]:
             `True` - if received status code 200 from Studio.
             `False` - if received other status code or RequestException raised.
             `None`- if prerequisites weren't met and the request was not sent.
     """
-    studio_token, studio_repo_url = get_studio_token_and_repo_url()
+    studio_token, studio_repo_url = get_studio_token_and_repo_url(
+        studio_token, studio_repo_url
+    )
 
     if any(x is None for x in (studio_token, studio_repo_url)):
         return None

diff --git a/tests/test_post_live_metrics.py b/tests/test_post_live_metrics.py
@@ -345,3 +345,36 @@ def test_get_studio_token_and_repo_url_skip_repo_url(monkeypatch):
     token, repo_url = get_studio_token_and_repo_url()
     assert token is None
     assert repo_url is None  # Skipped call to get_repo_url
+
+
+def test_post_live_metrics_token_and_repo_url_args(mocker, monkeypatch):
+    monkeypatch.setenv(DVC_STUDIO_URL, "https://0.0.0.0")
+
+    mocked_response = mocker.MagicMock()
+    mocked_response.status_code = 200
+    mocked_post = mocker.patch("requests.post", return_value=mocked_response)
+
+    assert post_live_metrics(
+        "start",
+        "f" * 40,
+        "fooname",
+        "fooclient",
+        studio_token="FOO_TOKEN",
+        studio_repo_url="FOO_REPO_URL",
+    )
+
+    mocked_post.assert_called_with(
+        "https://0.0.0.0/api/live",
+        json={
+            "type": "start",
+            "repo_url": "FOO_REPO_URL",
+            "baseline_sha": "f" * 40,
+            "name": "fooname",
+            "client": "fooclient",
+        },
+        headers={
+            "Authorization": "token FOO_TOKEN",
+            "Content-type": "application/json",
+        },
+        timeout=5,
+    )