From 125a3e269ae0a2d2ebc91a840381e1f954deca62 Mon Sep 17 00:00:00 2001 From: Douglas Blank Date: Fri, 13 Mar 2026 13:12:05 -0400 Subject: [PATCH 1/2] Fix review issues in migrate-users PR - copy.py: move `import requests` to top-level; use `self.api._get_url_server()` instead of broken `self.api.server_url`; split HTTPError vs RequestException handlers for precise error messages - migrate_users.py: add --url/--source-url args for self-hosted instances; fix _resolve_server_url (base64 padding, error handling, explicit-URL priority); require --source-api-key explicitly instead of silently reusing COMET_API_KEY; warn when source and dest URL+key are identical; handle dict-list API response in _get_existing_workspaces; show per-user output in dry-run; add --failures-output flag; use json= kwarg in _add_member; fix docstring examples - tests/unit/test_migrate_users.py: 17 new unit tests covering _resolve_server_url, _get_existing_workspaces, _add_member, and dry-run behaviour Co-Authored-By: Claude Sonnet 4.6 --- README.md | 22 ++-- cometx/cli/__init__.py | 1 + cometx/cli/copy.py | 16 ++- cometx/cli/migrate_users.py | 96 +++++++++++---- tests/unit/test_migrate_users.py | 205 +++++++++++++++++++++++++++++++ 5 files changed, 301 insertions(+), 39 deletions(-) create mode 100644 tests/unit/test_migrate_users.py diff --git a/README.md b/README.md index 3ced807..81ab650 100644 --- a/README.md +++ b/README.md @@ -333,9 +333,7 @@ For more information, `cometx log --help` ## cometx migrate-users -This command is used to migrate users into workspaces from a source Comet environment to a destination environment. It reads workspace membership from a chargeback report from the source environment and invites each user to the corresponding workspace in the destination environment by email. If a user with the email does not exist in the destination environment, they still will be provisioned access to these workspaces once they sign up using that email. - -The chargeback report can be fetched automatically from the source environment's admin API, or you can provide a pre-downloaded JSON file. +This command migrates users into workspaces from a source Comet environment to a destination environment. It reads workspace membership from a chargeback report fetched from the source environment's admin API (or a pre-downloaded local file) and invites each user to the corresponding workspace in the destination environment by email. If a user with that email does not yet exist in the destination, they will be provisioned access once they sign up. ``` cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY [FLAGS ...] @@ -343,14 +341,17 @@ cometx migrate-users --api-key DEST_KEY --chargeback-report /path/to/report.json ``` **Arguments:** -* `--api-key API_KEY` - API key for the destination environment where users will be added. Falls back to `COMET_API_KEY` environment variable if not provided. -* `--source-api-key SOURCE_API_KEY` - API key for the source environment. Used to fetch the chargeback report. Falls back to `COMET_API_KEY` environment variable if not provided. Required unless `--chargeback-report` is given. -* `--chargeback-report PATH` - Path to a local chargeback report JSON file. When provided, the report is loaded from this file instead of being fetched from the source environment, and `--source-api-key` is not required. +* `--api-key API_KEY` - API key for the destination environment. Falls back to `COMET_API_KEY` if not provided. +* `--url URL` - Base URL of the destination Comet environment (e.g. `https://comet.example.com`). Required for self-hosted instances when the API key does not encode the server URL. +* `--source-api-key SOURCE_API_KEY` - API key for the source environment. Used to fetch the chargeback report. Required unless `--chargeback-report` is given. +* `--source-url SOURCE_URL` - Base URL of the source Comet environment. Required for self-hosted source instances when the source API key does not encode the server URL. +* `--chargeback-report PATH` - Path to a local chargeback report JSON file. When provided, `--source-api-key` is not required. ### Flags -* `--create-workspaces` - Create workspaces on the destination environment if they don't already exist (default: off) +* `--create-workspaces` - Create workspaces on the destination if they don't already exist (default: off) * `--dry-run` - Print what would happen without making any changes +* `--failures-output PATH` - Path to write failed operations JSON (default: `bulk_add_failures_by_email.json`) ### Examples @@ -364,6 +365,9 @@ cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY # Use a local chargeback report file cometx migrate-users --api-key DEST_KEY --chargeback-report /tmp/chargeback_reports.json +# Self-hosted environments — provide explicit URLs +cometx migrate-users --api-key DEST_KEY --url https://comet.dest.example.com \ + --source-api-key SOURCE_KEY --source-url https://comet.src.example.com ``` For more information, `cometx migrate-users --help` @@ -518,12 +522,8 @@ cometx admin chargeback-report [YEAR-MONTH] **Examples:** ``` cometx admin chargeback-report -<<<<<<< Updated upstream -cometx admin usage-report -======= cometx admin chargeback-report 2024-09 ``` ->>>>>>> Stashed changes #### usage-report diff --git a/cometx/cli/__init__.py b/cometx/cli/__init__.py index 975b9d1..067e86c 100644 --- a/cometx/cli/__init__.py +++ b/cometx/cli/__init__.py @@ -108,6 +108,7 @@ def main(raw_args=sys.argv[1:]): add_subparser(subparsers, count, "count") add_subparser(subparsers, reproduce, "reproduce") add_subparser(subparsers, config, "config") + add_subparser(subparsers, migrate_users, "migrate-users") add_subparser(subparsers, rename_duplicates, "rename-duplicates") add_subparser(subparsers, smoke_test, "smoke-test") add_subparser(subparsers, migrate_users, "migrate-users") diff --git a/cometx/cli/copy.py b/cometx/cli/copy.py index 5a21c82..ac60b78 100644 --- a/cometx/cli/copy.py +++ b/cometx/cli/copy.py @@ -78,6 +78,8 @@ import time import urllib.parse import zipfile + +import requests from datetime import datetime, timedelta from comet_ml import APIExperiment, Artifact, Experiment, OfflineExperiment @@ -834,10 +836,8 @@ def copy(self, source, destination, symlink, ignore, debug, sync, create_workspa print( f"Workspace {workspace_dst!r} does not exist, attempting to create it..." ) + url = f"{self.api._get_url_server()}/api/rest/v2/write/workspace/new" try: - import requests - - url = f"{self.api.server_url}/api/rest/v2/write/workspace/new" response = requests.post( url, json={"name": workspace_dst}, @@ -848,12 +848,18 @@ def copy(self, source, destination, symlink, ignore, debug, sync, create_workspa ) response.raise_for_status() print(f"Workspace {workspace_dst!r} created successfully.") - except Exception as exc: + except requests.exceptions.HTTPError as exc: + raise Exception( + f"Workspace {workspace_dst!r} does not exist and could not be " + f"created automatically (HTTP {exc.response.status_code}). " + f"Please create it via the Comet UI and try again." + ) from exc + except requests.exceptions.RequestException as exc: raise Exception( f"Workspace {workspace_dst!r} does not exist and could not be " f"created automatically: {exc}. " f"Please create it via the Comet UI and try again." - ) + ) from exc else: raise Exception( f"{workspace_dst} does not exist; use --create-workspaces to " diff --git a/cometx/cli/migrate_users.py b/cometx/cli/migrate_users.py index e7997de..67fca5a 100644 --- a/cometx/cli/migrate_users.py +++ b/cometx/cli/migrate_users.py @@ -20,9 +20,9 @@ Examples: -$ cometx --api-key DEST_KEY migrate-users --source-api-key SOURCE_KEY --dry-run -$ cometx --api-key DEST_KEY migrate-users --source-api-key SOURCE_KEY -$ cometx --api-key DEST_KEY migrate-users --chargeback-report /path/to/report.json +$ cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY --dry-run +$ cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY +$ cometx migrate-users --api-key DEST_KEY --chargeback-report /path/to/report.json """ import argparse @@ -30,6 +30,7 @@ import json import os import sys + import requests ADDITIONAL_ARGS = False @@ -40,13 +41,29 @@ def get_parser_arguments(parser): parser.add_argument( "--api-key", - help="API key for the destination environment (used to add workspace members)", + help="API key for the destination environment (used to add workspace members). " + "Falls back to COMET_API_KEY environment variable if not provided.", + type=str, + default=None, + ) + parser.add_argument( + "--url", + help="Base URL of the destination Comet environment (e.g. https://comet.example.com). " + "Required for self-hosted instances when the API key does not encode the server URL.", type=str, default=None, ) parser.add_argument( "--source-api-key", - help="API key for the source environment (used to fetch the chargeback report)", + help="API key for the source environment (used to fetch the chargeback report). " + "Required unless --chargeback-report is given.", + type=str, + default=None, + ) + parser.add_argument( + "--source-url", + help="Base URL of the source Comet environment. " + "Required for self-hosted instances when the source API key does not encode the server URL.", type=str, default=None, ) @@ -68,15 +85,36 @@ def get_parser_arguments(parser): default=False, action="store_true", ) + parser.add_argument( + "--failures-output", + help="Path to write failed operations JSON (default: bulk_add_failures_by_email.json)", + type=str, + default="bulk_add_failures_by_email.json", + ) -def _resolve_server_url(api_key): - if "*" not in api_key: - return COMET_CLOUD_URL +def _resolve_server_url(api_key, explicit_url=None): + """Return the server base URL. - _, encoded = api_key.split("*", 1) - payload = json.loads(base64.b64decode(encoded)) - return payload["baseUrl"].rstrip("/") + Priority: + 1. ``explicit_url`` (from --url / --source-url), if provided. + 2. URL encoded inside the API key (new-style keys contain ``*``). + 3. COMET_CLOUD_URL as a last resort. + """ + if explicit_url: + return explicit_url.rstrip("/") + + if "*" in api_key: + try: + _, encoded = api_key.split("*", 1) + # base64 padding may be missing + padding = (4 - len(encoded) % 4) % 4 + payload = json.loads(base64.b64decode(encoded + "=" * padding)) + return payload["baseUrl"].rstrip("/") + except Exception: + pass # Fall through to default + + return COMET_CLOUD_URL def _fetch_chargeback_report(server_url, source_api_key): @@ -101,7 +139,11 @@ def _get_existing_workspaces(dest_url, headers): url = f"{dest_url}/api/rest/v2/workspaces" resp = requests.get(url, headers=headers, timeout=15) resp.raise_for_status() - return set(resp.json()) + data = resp.json() + # The endpoint may return a list of strings or a list of dicts with a "name" key. + if data and isinstance(data[0], dict): + return {ws["name"] for ws in data} + return set(data) def _create_workspace(dest_url, headers, workspace_name): @@ -126,7 +168,7 @@ def _add_member(url, headers, email, workspace_name): response = requests.post( url, headers=headers, - data=json.dumps(payload), + json=payload, timeout=15, ) @@ -163,22 +205,30 @@ def migrate_users(parsed_args): print("[ERROR] No API key found. Set COMET_API_KEY or pass --api-key.") sys.exit(1) - source_api_key = parsed_args.source_api_key or os.environ.get("COMET_API_KEY") + source_api_key = parsed_args.source_api_key create_workspaces = parsed_args.create_workspaces dry_run = parsed_args.dry_run + failures_output = parsed_args.failures_output if not parsed_args.chargeback_report and not source_api_key: - print("[ERROR] Provide either --source-api-key or --chargeback-report.") + print( + "[ERROR] --source-api-key is required when --chargeback-report is not provided." + ) sys.exit(1) - dest_url = _resolve_server_url(api_key) + dest_url = _resolve_server_url(api_key, parsed_args.url) print(f"Destination URL: {dest_url}") if parsed_args.chargeback_report: data = _load_chargeback_report(parsed_args.chargeback_report) else: - source_url = _resolve_server_url(source_api_key) + source_url = _resolve_server_url(source_api_key, parsed_args.source_url) print(f"Source URL: {source_url}") + if source_url == dest_url and source_api_key == api_key: + print( + "[WARNING] Source and destination URL and API key are identical. " + "Are you sure you want to migrate users to the same environment?" + ) try: data = _fetch_chargeback_report(source_url, source_api_key) except requests.exceptions.RequestException as e: @@ -215,7 +265,7 @@ def migrate_users(parsed_args): elif create_workspaces and dry_run: print("[DRY RUN] Would check/create workspaces on destination\n") - url = f"{dest_url}/api/rest/v2/write/add-workspace-member" + add_member_url = f"{dest_url}/api/rest/v2/write/add-workspace-member" total_added = 0 total_skipped = 0 @@ -252,11 +302,12 @@ def migrate_users(parsed_args): continue if dry_run: + print(f" [DRY RUN] Would add '{email}' to '{ws_name}'") ws_added += 1 total_added += 1 continue - status, error_info = _add_member(url, headers, email, ws_name) + status, error_info = _add_member(add_member_url, headers, email, ws_name) if status == "added": ws_added += 1 @@ -280,7 +331,7 @@ def migrate_users(parsed_args): ) if dry_run: - print(f" [DRY RUN] Would add {ws_added}/{len(members)} users") + print(f" [DRY RUN] Total: would add {ws_added}/{len(members)} users") else: parts = [f" Added {ws_added}/{len(members)} users successfully"] if ws_already_member: @@ -309,10 +360,9 @@ def migrate_users(parsed_args): print(f" *** Remove --dry-run to execute for real ***") if failures: - failures_file = "bulk_add_failures_by_email.json" - with open(failures_file, "w") as f: + with open(failures_output, "w") as f: json.dump(failures, f, indent=2) - print(f"\n Failed operations saved to {failures_file}") + print(f"\n Failed operations saved to {failures_output}") def main(args): diff --git a/tests/unit/test_migrate_users.py b/tests/unit/test_migrate_users.py new file mode 100644 index 0000000..3547b00 --- /dev/null +++ b/tests/unit/test_migrate_users.py @@ -0,0 +1,205 @@ +# -*- coding: utf-8 -*- +# **************************************** +# __ +# _________ ____ ___ ___ / /__ __ +# / ___/ __ \/ __ `__ \/ _ \/ __/ |/_/ +# / /__/ /_/ / / / / / / __/ /__> < +# \___/\____/_/ /_/ /_/\___/\__/_/|_| +# +# +# Copyright (c) 2022 Cometx Development +# Team. All rights reserved. +# **************************************** + +import argparse +import base64 +import json +import unittest +from unittest.mock import MagicMock, call, patch + +from cometx.cli.migrate_users import ( + COMET_CLOUD_URL, + _add_member, + _get_existing_workspaces, + _resolve_server_url, + migrate_users, +) + + +def _make_new_style_key(base_url): + payload = json.dumps({"baseUrl": base_url}).encode() + encoded = base64.b64encode(payload).decode().rstrip("=") + return f"abc123*{encoded}" + + +class TestResolveServerUrl: + def test_explicit_url_takes_priority(self): + key = _make_new_style_key("https://encoded.example.com") + assert _resolve_server_url(key, "https://explicit.example.com") == "https://explicit.example.com" + + def test_explicit_url_trailing_slash_stripped(self): + assert _resolve_server_url("anykey", "https://example.com/") == "https://example.com" + + def test_new_style_key_decoded(self): + key = _make_new_style_key("https://self-hosted.example.com") + assert _resolve_server_url(key) == "https://self-hosted.example.com" + + def test_new_style_key_trailing_slash_stripped(self): + key = _make_new_style_key("https://self-hosted.example.com/") + assert _resolve_server_url(key) == "https://self-hosted.example.com" + + def test_old_style_key_defaults_to_cloud(self): + assert _resolve_server_url("oldstylekey") == COMET_CLOUD_URL + + def test_malformed_new_style_key_defaults_to_cloud(self): + assert _resolve_server_url("abc*notvalidbase64!!!") == COMET_CLOUD_URL + + def test_new_style_key_missing_base_url_defaults_to_cloud(self): + payload = json.dumps({"other": "value"}).encode() + encoded = base64.b64encode(payload).decode() + key = f"abc*{encoded}" + assert _resolve_server_url(key) == COMET_CLOUD_URL + + +class TestGetExistingWorkspaces: + @patch("cometx.cli.migrate_users.requests.get") + def test_flat_list_response(self, mock_get): + mock_get.return_value.json.return_value = ["ws1", "ws2"] + mock_get.return_value.raise_for_status = MagicMock() + result = _get_existing_workspaces("https://example.com", {}) + assert result == {"ws1", "ws2"} + + @patch("cometx.cli.migrate_users.requests.get") + def test_dict_list_response(self, mock_get): + mock_get.return_value.json.return_value = [{"name": "ws1"}, {"name": "ws2"}] + mock_get.return_value.raise_for_status = MagicMock() + result = _get_existing_workspaces("https://example.com", {}) + assert result == {"ws1", "ws2"} + + +class TestAddMember: + @patch("cometx.cli.migrate_users.requests.post") + def test_success(self, mock_post): + mock_post.return_value.status_code = 200 + status, error = _add_member("http://example.com/add", {}, "user@example.com", "ws1") + assert status == "added" + assert error is None + + @patch("cometx.cli.migrate_users.requests.post") + def test_already_member(self, mock_post): + mock_post.return_value.status_code = 400 + mock_post.return_value.json.return_value = {"msg": "already member of ws1"} + status, error = _add_member("http://example.com/add", {}, "user@example.com", "ws1") + assert status == "already_member" + assert error is None + + @patch("cometx.cli.migrate_users.requests.post") + def test_failure(self, mock_post): + mock_post.return_value.status_code = 403 + mock_post.return_value.json.return_value = {"msg": "forbidden"} + status, error = _add_member("http://example.com/add", {}, "user@example.com", "ws1") + assert status == "failed" + assert error["status"] == 403 + + @patch("cometx.cli.migrate_users.requests.post") + def test_request_exception(self, mock_post): + import requests as req + mock_post.side_effect = req.exceptions.ConnectionError("timeout") + status, error = _add_member("http://example.com/add", {}, "user@example.com", "ws1") + assert status == "failed" + assert error["status"] == "exception" + + +class TestMigrateUsersDryRun: + def _make_args(self, **kwargs): + defaults = dict( + api_key="dest-key", + url=None, + source_api_key="source-key", + source_url=None, + chargeback_report=None, + create_workspaces=False, + dry_run=True, + failures_output="failures.json", + ) + defaults.update(kwargs) + return argparse.Namespace(**defaults) + + @patch("cometx.cli.migrate_users._resolve_server_url", return_value="https://dest.example.com") + @patch("cometx.cli.migrate_users._fetch_chargeback_report") + @patch("builtins.print") + def test_dry_run_no_api_calls(self, mock_print, mock_fetch, mock_resolve): + mock_fetch.return_value = { + "workspaces": [ + { + "name": "ws1", + "members": [ + {"email": "alice@example.com", "userName": "alice"}, + {"email": "bob@example.com", "userName": "bob"}, + ], + } + ] + } + with patch("cometx.cli.migrate_users.requests.post") as mock_post: + migrate_users(self._make_args()) + mock_post.assert_not_called() + + @patch("cometx.cli.migrate_users._resolve_server_url", return_value="https://dest.example.com") + @patch("cometx.cli.migrate_users._fetch_chargeback_report") + @patch("builtins.print") + def test_dry_run_prints_per_user(self, mock_print, mock_fetch, mock_resolve): + mock_fetch.return_value = { + "workspaces": [ + { + "name": "ws1", + "members": [ + {"email": "alice@example.com", "userName": "alice"}, + ], + } + ] + } + migrate_users(self._make_args()) + printed = " ".join(str(c) for c in mock_print.call_args_list) + assert "alice@example.com" in printed + assert "ws1" in printed + + @patch("cometx.cli.migrate_users._resolve_server_url", return_value="https://dest.example.com") + @patch("cometx.cli.migrate_users._fetch_chargeback_report") + @patch("builtins.print") + def test_dry_run_skips_no_email(self, mock_print, mock_fetch, mock_resolve): + mock_fetch.return_value = { + "workspaces": [ + { + "name": "ws1", + "members": [ + {"email": None, "userName": "noemail"}, + ], + } + ] + } + migrate_users(self._make_args()) + printed = " ".join(str(c) for c in mock_print.call_args_list) + assert "noemail" in printed + assert "no email" in printed + + def test_no_source_key_without_chargeback_report_exits(self): + args = self._make_args(source_api_key=None, chargeback_report=None) + with self.assertRaises(SystemExit): + migrate_users(args) + + def assertRaises(self, exc): + import contextlib + + @contextlib.contextmanager + def ctx(): + try: + yield + raise AssertionError(f"{exc} was not raised") + except exc: + pass + + return ctx() + + +if __name__ == "__main__": + unittest.main() From 518bbaedd094ef7b48467ee0a51dc7d850e6fe6b Mon Sep 17 00:00:00 2001 From: Douglas Blank Date: Fri, 13 Mar 2026 13:56:42 -0400 Subject: [PATCH 2/2] Address PR #21 review comments - Fix duplicate migrate-users subparser registration in __init__.py - Validate explicit URLs require https:// in _resolve_server_url (SSRF fix) - Remove silent COMET_CLOUD_URL fallback; fail fast when URL cannot be determined - Handle {"workspaceNames": [...]} dict response shape in _get_existing_workspaces - Deduplicate workspace creation by reusing _create_workspace from migrate_users in copy.py - Update tests to match new behavior (SystemExit instead of COMET_CLOUD_URL fallback, new dict response shapes) Co-Authored-By: Claude Sonnet 4.6 --- cometx/cli/__init__.py | 1 - cometx/cli/copy.py | 17 ++++++-------- cometx/cli/migrate_users.py | 25 +++++++++++++++++---- tests/unit/test_migrate_users.py | 38 ++++++++++++++++++++++++++------ 4 files changed, 59 insertions(+), 22 deletions(-) diff --git a/cometx/cli/__init__.py b/cometx/cli/__init__.py index 067e86c..975b9d1 100644 --- a/cometx/cli/__init__.py +++ b/cometx/cli/__init__.py @@ -108,7 +108,6 @@ def main(raw_args=sys.argv[1:]): add_subparser(subparsers, count, "count") add_subparser(subparsers, reproduce, "reproduce") add_subparser(subparsers, config, "config") - add_subparser(subparsers, migrate_users, "migrate-users") add_subparser(subparsers, rename_duplicates, "rename-duplicates") add_subparser(subparsers, smoke_test, "smoke-test") add_subparser(subparsers, migrate_users, "migrate-users") diff --git a/cometx/cli/copy.py b/cometx/cli/copy.py index ac60b78..11357ad 100644 --- a/cometx/cli/copy.py +++ b/cometx/cli/copy.py @@ -104,6 +104,7 @@ from ..framework.comet.download_manager import sanitize_filename from ..utils import remove_extra_slashes from .copy_utils import upload_single_offline_experiment +from .migrate_users import _create_workspace ADDITIONAL_ARGS = False @@ -836,17 +837,13 @@ def copy(self, source, destination, symlink, ignore, debug, sync, create_workspa print( f"Workspace {workspace_dst!r} does not exist, attempting to create it..." ) - url = f"{self.api._get_url_server()}/api/rest/v2/write/workspace/new" + dest_url = self.api._get_url_server() + headers = { + "Authorization": self.api.api_key, + "Content-Type": "application/json", + } try: - response = requests.post( - url, - json={"name": workspace_dst}, - headers={ - "Authorization": self.api.api_key, - "Content-Type": "application/json", - }, - ) - response.raise_for_status() + _create_workspace(dest_url, headers, workspace_dst) print(f"Workspace {workspace_dst!r} created successfully.") except requests.exceptions.HTTPError as exc: raise Exception( diff --git a/cometx/cli/migrate_users.py b/cometx/cli/migrate_users.py index 67fca5a..ed7841d 100644 --- a/cometx/cli/migrate_users.py +++ b/cometx/cli/migrate_users.py @@ -30,6 +30,7 @@ import json import os import sys +import urllib.parse import requests @@ -99,9 +100,13 @@ def _resolve_server_url(api_key, explicit_url=None): Priority: 1. ``explicit_url`` (from --url / --source-url), if provided. 2. URL encoded inside the API key (new-style keys contain ``*``). - 3. COMET_CLOUD_URL as a last resort. + 3. Error — no silent fallback to cloud URL. """ if explicit_url: + parsed = urllib.parse.urlparse(explicit_url) + if parsed.scheme != "https": + print("[ERROR] --url/--source-url must use https://.") + sys.exit(1) return explicit_url.rstrip("/") if "*" in api_key: @@ -112,9 +117,14 @@ def _resolve_server_url(api_key, explicit_url=None): payload = json.loads(base64.b64decode(encoded + "=" * padding)) return payload["baseUrl"].rstrip("/") except Exception: - pass # Fall through to default + pass # Fall through to error below - return COMET_CLOUD_URL + print( + "[ERROR] Cannot determine the server URL from the API key. " + "Pass --url (destination) or --source-url (source) explicitly, " + "e.g. --url https://comet.example.com" + ) + sys.exit(1) def _fetch_chargeback_report(server_url, source_api_key): @@ -140,9 +150,16 @@ def _get_existing_workspaces(dest_url, headers): resp = requests.get(url, headers=headers, timeout=15) resp.raise_for_status() data = resp.json() - # The endpoint may return a list of strings or a list of dicts with a "name" key. + # Handle {"workspaceNames": [...]} dict shape + if isinstance(data, dict): + names = data.get("workspaceNames", []) + if names and isinstance(names[0], dict): + return {ws["name"] for ws in names} + return set(names) + # Handle list of dicts with a "name" key if data and isinstance(data[0], dict): return {ws["name"] for ws in data} + # Handle flat list of strings return set(data) diff --git a/tests/unit/test_migrate_users.py b/tests/unit/test_migrate_users.py index 3547b00..e812ee9 100644 --- a/tests/unit/test_migrate_users.py +++ b/tests/unit/test_migrate_users.py @@ -17,8 +17,9 @@ import unittest from unittest.mock import MagicMock, call, patch +import pytest + from cometx.cli.migrate_users import ( - COMET_CLOUD_URL, _add_member, _get_existing_workspaces, _resolve_server_url, @@ -40,6 +41,10 @@ def test_explicit_url_takes_priority(self): def test_explicit_url_trailing_slash_stripped(self): assert _resolve_server_url("anykey", "https://example.com/") == "https://example.com" + def test_explicit_url_non_https_exits(self): + with pytest.raises(SystemExit): + _resolve_server_url("anykey", "http://example.com") + def test_new_style_key_decoded(self): key = _make_new_style_key("https://self-hosted.example.com") assert _resolve_server_url(key) == "https://self-hosted.example.com" @@ -48,17 +53,20 @@ def test_new_style_key_trailing_slash_stripped(self): key = _make_new_style_key("https://self-hosted.example.com/") assert _resolve_server_url(key) == "https://self-hosted.example.com" - def test_old_style_key_defaults_to_cloud(self): - assert _resolve_server_url("oldstylekey") == COMET_CLOUD_URL + def test_old_style_key_exits(self): + with pytest.raises(SystemExit): + _resolve_server_url("oldstylekey") - def test_malformed_new_style_key_defaults_to_cloud(self): - assert _resolve_server_url("abc*notvalidbase64!!!") == COMET_CLOUD_URL + def test_malformed_new_style_key_exits(self): + with pytest.raises(SystemExit): + _resolve_server_url("abc*notvalidbase64!!!") - def test_new_style_key_missing_base_url_defaults_to_cloud(self): + def test_new_style_key_missing_base_url_exits(self): payload = json.dumps({"other": "value"}).encode() encoded = base64.b64encode(payload).decode() key = f"abc*{encoded}" - assert _resolve_server_url(key) == COMET_CLOUD_URL + with pytest.raises(SystemExit): + _resolve_server_url(key) class TestGetExistingWorkspaces: @@ -76,6 +84,22 @@ def test_dict_list_response(self, mock_get): result = _get_existing_workspaces("https://example.com", {}) assert result == {"ws1", "ws2"} + @patch("cometx.cli.migrate_users.requests.get") + def test_workspace_names_dict_response(self, mock_get): + mock_get.return_value.json.return_value = {"workspaceNames": ["ws1", "ws2"]} + mock_get.return_value.raise_for_status = MagicMock() + result = _get_existing_workspaces("https://example.com", {}) + assert result == {"ws1", "ws2"} + + @patch("cometx.cli.migrate_users.requests.get") + def test_workspace_names_dict_with_dict_entries(self, mock_get): + mock_get.return_value.json.return_value = { + "workspaceNames": [{"name": "ws1"}, {"name": "ws2"}] + } + mock_get.return_value.raise_for_status = MagicMock() + result = _get_existing_workspaces("https://example.com", {}) + assert result == {"ws1", "ws2"} + class TestAddMember: @patch("cometx.cli.migrate_users.requests.post")