Support creating repositories under organizations during init

.gitignore

@@ -126,6 +126,7 @@ dmypy.json
 
 # IntelliJ
 .idea/
+*.iml
 
 # VSCode
 .vscode/

dagshub/common/init.py

@@ -87,16 +87,27 @@ def init(
     try:
         repo_api.get_repo_info()
     except RepoNotFoundError:
-        log_message(f"Repository {repo_name} doesn't exist, creating it under current user.")
-        create_repo(repo_name)
+        should_create_under_org = False
+        if repo_owner:
+            current_user = UserAPI.get_current_user(host=host)
+            should_create_under_org = repo_owner != current_user.username
+
+        if should_create_under_org:
+            log_message(
+                f'Repository {repo_name} doesn\'t exist, creating it under organization "{repo_owner}".'
+            )
+            create_repo(repo_name, org_name=repo_owner, host=host)
+        else:
+            log_message(f"Repository {repo_name} doesn't exist, creating it under current user.")
+            create_repo(repo_name, host=host)
 
     # Get the token for the configs
     token = get_token(host=host)
 
     # Configure MLFlow
     if mlflow:
         os.environ["MLFLOW_TRACKING_URI"] = f"{url}.mlflow"
-        os.environ["MLFLOW_TRACKING_USERNAME"] = UserAPI.get_user_from_token(token).username
+        os.environ["MLFLOW_TRACKING_USERNAME"] = UserAPI.get_user_from_token(token, host=host).username
         os.environ["MLFLOW_TRACKING_PASSWORD"] = token
 
         log_message(f'Initialized MLflow to track repo "{repo_owner}/{repo_name}"')

tests/common/test_init.py

@@ -0,0 +1,86 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+import dagshub
+from dagshub.common.api.repo import RepoNotFoundError
+
+
+@pytest.fixture
+def mock_repo_api(mocker):
+    mock = mocker.patch("dagshub.common.init.RepoAPI")
+    mock.return_value.get_repo_info.side_effect = RepoNotFoundError()
+    return mock
+
+
+@pytest.fixture
+def mock_user_api(mocker):
+    user = MagicMock()
+    user.username = "testuser"
+    mock = mocker.patch("dagshub.common.init.UserAPI")
+    mock.get_current_user.return_value = user
+    mock.get_user_from_token.return_value = user
+    return mock
+
+
+@pytest.fixture
+def mock_create_repo(mocker):
+    return mocker.patch("dagshub.common.init.create_repo")
+
+
+@pytest.fixture
+def mock_get_token(mocker):
+    return mocker.patch("dagshub.common.init.get_token", return_value="fake-token")
+
+
+@pytest.fixture
+def mock_log_message(mocker):
+    return mocker.patch("dagshub.common.init.log_message")
+
+
+def test_init_creates_repo_under_org_when_owner_differs(
+    mock_repo_api, mock_user_api, mock_create_repo, mock_get_token, mock_log_message
+):
+    dagshub.init(repo_owner="my-org", repo_name="my-repo", mlflow=False, dvc=False)
+
+    mock_user_api.get_current_user.assert_called_once()
+    mock_create_repo.assert_called_once_with("my-repo", org_name="my-org", host="https://dagshub.com")
+    mock_log_message.assert_any_call(
+        'Repository my-repo doesn\'t exist, creating it under organization "my-org".'
+    )
+
+
+def test_init_creates_repo_under_current_user_when_owner_matches(
+    mock_repo_api, mock_user_api, mock_create_repo, mock_get_token, mock_log_message
+):
+    dagshub.init(repo_owner="testuser", repo_name="my-repo", mlflow=False, dvc=False)
+
+    mock_user_api.get_current_user.assert_called_once()
+    mock_create_repo.assert_called_once_with("my-repo", host="https://dagshub.com")
+    mock_log_message.assert_any_call(
+        "Repository my-repo doesn't exist, creating it under current user."
+    )
+
+
+def test_init_creates_repo_under_current_user_from_url(
+    mock_repo_api, mock_user_api, mock_create_repo, mock_get_token, mock_log_message
+):
+    dagshub.init(url="https://dagshub.com/testuser/my-repo", mlflow=False, dvc=False)
+
+    mock_user_api.get_current_user.assert_called_once()
+    mock_create_repo.assert_called_once_with("my-repo", host="https://dagshub.com")
+    mock_log_message.assert_any_call(
+        "Repository my-repo doesn't exist, creating it under current user."
+    )
+
+
+def test_init_creates_repo_under_org_from_url(
+    mock_repo_api, mock_user_api, mock_create_repo, mock_get_token, mock_log_message
+):
+    dagshub.init(url="https://dagshub.com/my-org/my-repo", mlflow=False, dvc=False)
+
+    mock_user_api.get_current_user.assert_called_once()
+    mock_create_repo.assert_called_once_with("my-repo", org_name="my-org", host="https://dagshub.com")
+    mock_log_message.assert_any_call(
+        'Repository my-repo doesn\'t exist, creating it under organization "my-org".'
+    )