datahub/smoke-test/tests/cli/user_cmd/test_user_add.py

import logging
import uuid
from typing import Any, List

import pytest

from datahub.ingestion.graph.client import DataHubGraph
from tests.consistency_utils import wait_for_writes_to_sync
from tests.utils import run_datahub_cmd

logger = logging.getLogger(__name__)


def generate_test_email() -> str:
    """Generate a unique email for testing to avoid conflicts."""
    return f"test-user-{uuid.uuid4()}@example.com"


def datahub_user_add(
    auth_session: Any,
    email: str,
    display_name: str,
    password: str,
    role: str | None = None,
    user_id: str | None = None,
    email_as_id: bool = False,
) -> Any:
    """Run the datahub user add command."""
    add_args: List[str] = [
        "user",
        "add",
        "--email",
        email,
        "--display-name",
        display_name,
        "--password",
    ]

    if user_id:
        add_args.extend(["--id", user_id])
    elif email_as_id:
        add_args.append("--email-as-id")

    if role:
        add_args.extend(["--role", role])

    result = run_datahub_cmd(
        add_args,
        input=f"{password}\n{password}\n",
        env={
            "DATAHUB_GMS_URL": auth_session.gms_url(),
            "DATAHUB_GMS_TOKEN": auth_session.gms_token(),
        },
    )
    return result


@pytest.fixture
def test_users_cleanup(auth_session: Any, graph_client: DataHubGraph):
    """Fixture to track and clean up test users."""
    created_users: List[str] = []

    def add_user_to_cleanup(email: str):
        user_urn = f"urn:li:corpuser:{email}"
        created_users.append(user_urn)

    yield add_user_to_cleanup

    for user_urn in created_users:
        try:
            graph_client.hard_delete_entity(user_urn)
        except Exception as e:
            logger.warning(f"Failed to clean up user {user_urn}: {e}")

    if created_users:
        wait_for_writes_to_sync()


def test_user_add_without_role(auth_session: Any, test_users_cleanup: Any) -> None:
    """Test creating a user without specifying a role."""
    email = generate_test_email()
    display_name = "Test User No Role"
    password = "testpassword123"

    test_users_cleanup(email)

    result = datahub_user_add(
        auth_session, email, display_name, password, email_as_id=True
    )

    assert result.exit_code == 0
    assert "Successfully created user" in result.output
    assert email in result.output
    assert "URN:" in result.output

    wait_for_writes_to_sync()


def test_user_add_with_admin_role(auth_session: Any, test_users_cleanup: Any) -> None:
    """Test creating a user with Admin role."""
    email = generate_test_email()
    display_name = "Test Admin User"
    password = "adminpass123"

    test_users_cleanup(email)

    result = datahub_user_add(
        auth_session, email, display_name, password, role="Admin", email_as_id=True
    )

    assert result.exit_code == 0
    assert "Successfully created user" in result.output
    assert email in result.output
    assert "Admin" in result.output
    assert "URN:" in result.output

    wait_for_writes_to_sync()


def test_user_add_with_editor_role(auth_session: Any, test_users_cleanup: Any) -> None:
    """Test creating a user with Editor role."""
    email = generate_test_email()
    display_name = "Test Editor User"
    password = "editorpass123"

    test_users_cleanup(email)

    result = datahub_user_add(
        auth_session, email, display_name, password, role="editor", email_as_id=True
    )

    assert result.exit_code == 0
    assert "Successfully created user" in result.output
    assert email in result.output
    assert "Editor" in result.output
    assert "URN:" in result.output

    wait_for_writes_to_sync()


def test_user_add_with_reader_role(auth_session: Any, test_users_cleanup: Any) -> None:
    """Test creating a user with Reader role."""
    email = generate_test_email()
    display_name = "Test Reader User"
    password = "readerpass123"

    test_users_cleanup(email)

    result = datahub_user_add(
        auth_session, email, display_name, password, role="READER", email_as_id=True
    )

    assert result.exit_code == 0
    assert "Successfully created user" in result.output
    assert email in result.output
    assert "Reader" in result.output
    assert "URN:" in result.output

    wait_for_writes_to_sync()


def test_user_add_duplicate_user(auth_session: Any, test_users_cleanup: Any) -> None:
    """Test that creating a duplicate user returns an error."""
    email = generate_test_email()
    display_name = "Test Duplicate User"
    password = "duplicatepass123"

    test_users_cleanup(email)

    first_result = datahub_user_add(
        auth_session, email, display_name, password, email_as_id=True
    )
    assert first_result.exit_code == 0

    wait_for_writes_to_sync()

    second_result = datahub_user_add(
        auth_session, email, display_name, password, email_as_id=True
    )
    assert second_result.exit_code == 1
    assert "already exists" in second_result.output.lower()
    assert email in second_result.output


def test_user_add_without_password_flag(auth_session: Any) -> None:
    """Test that not providing --password flag results in error."""
    email = generate_test_email()
    display_name = "Test User No Password Flag"

    add_args: List[str] = [
        "user",
        "add",
        "--email",
        email,
        "--display-name",
        display_name,
        "--email-as-id",
    ]

    result = run_datahub_cmd(
        add_args,
        env={
            "DATAHUB_GMS_URL": auth_session.gms_url(),
            "DATAHUB_GMS_TOKEN": auth_session.gms_token(),
        },
    )

    assert result.exit_code == 1
    assert "password flag is required" in result.output.lower()


def test_user_add_with_explicit_id(auth_session: Any, test_users_cleanup: Any) -> None:
    """Test creating a user with explicit --id option."""
    email = generate_test_email()
    user_id = f"testuser_{uuid.uuid4().hex[:8]}"
    display_name = "Test User With ID"
    password = "testpassword123"

    # Cleanup with the user ID, not email
    test_users_cleanup(user_id)

    result = datahub_user_add(
        auth_session, email, display_name, password, user_id=user_id
    )

    assert result.exit_code == 0
    assert "Successfully created user" in result.output
    assert user_id in result.output
    assert "URN:" in result.output

    wait_for_writes_to_sync()


def test_user_add_without_id_or_flag(auth_session: Any) -> None:
    """Test that not providing --id or --email-as-id results in error."""
    email = generate_test_email()
    display_name = "Test User No ID"

    add_args: List[str] = [
        "user",
        "add",
        "--email",
        email,
        "--display-name",
        display_name,
        "--password",
    ]

    result = run_datahub_cmd(
        add_args,
        input="password123\npassword123\n",
        env={
            "DATAHUB_GMS_URL": auth_session.gms_url(),
            "DATAHUB_GMS_TOKEN": auth_session.gms_token(),
        },
    )

    assert result.exit_code == 1
    assert "must specify either --id or --email-as-id" in result.output.lower()


def test_user_add_with_both_id_and_flag(auth_session: Any) -> None:
    """Test that providing both --id and --email-as-id results in error."""
    email = generate_test_email()
    display_name = "Test User Both Options"

    add_args: List[str] = [
        "user",
        "add",
        "--email",
        email,
        "--id",
        "testuser",
        "--email-as-id",
        "--display-name",
        display_name,
        "--password",
    ]

    result = run_datahub_cmd(
        add_args,
        input="password123\npassword123\n",
        env={
            "DATAHUB_GMS_URL": auth_session.gms_url(),
            "DATAHUB_GMS_TOKEN": auth_session.gms_token(),
        },
    )

    assert result.exit_code == 1
    assert "cannot specify both --id and --email-as-id" in result.output.lower()


def test_user_add_id_different_from_email(
    auth_session: Any, test_users_cleanup: Any
) -> None:
    """Test creating a user where ID is different from email."""
    email = "john.doe@company.com"
    user_id = "jdoe"
    display_name = "John Doe"
    password = "securepass123"

    test_users_cleanup(user_id)

    result = datahub_user_add(
        auth_session, email, display_name, password, user_id=user_id, role="Editor"
    )

    assert result.exit_code == 0
    assert "Successfully created user" in result.output
    assert user_id in result.output
    assert "Editor" in result.output
    assert f"urn:li:corpuser:{user_id}" in result.output

    wait_for_writes_to_sync()