OpenMetadata/ingestion/tests/unit/sdk/test_user_entity.py

"""
Comprehensive unit tests for User entity with full mock coverage.
"""
import unittest
from unittest.mock import MagicMock
from uuid import UUID

from metadata.generated.schema.api.teams.createUser import CreateUserRequest
from metadata.generated.schema.entity.teams.user import AuthenticationMechanism
from metadata.generated.schema.entity.teams.user import User as UserEntity
from metadata.generated.schema.type.entityReference import EntityReference
from metadata.sdk import Users


class TestUserEntity(unittest.TestCase):
    """Comprehensive tests for User entity operations"""

    def setUp(self):
        """Set up test fixtures"""
        self.mock_ometa = MagicMock()

        # Set default client directly
        Users.set_default_client(self.mock_ometa)

        # Test data
        self.user_id = "350e8400-e29b-41d4-a716-446655440000"
        self.user_fqn = "john.doe"

    def test_create_user(self):
        """Test creating a user"""
        # Arrange
        create_request = CreateUserRequest(
            name="john.doe",
            email="john.doe@company.com",
            displayName="John Doe",
            description="Senior Data Engineer",
        )

        expected_user = MagicMock(spec=UserEntity)
        expected_user.id = UUID(self.user_id)
        expected_user.name = "john.doe"
        expected_user.email = "john.doe@company.com"
        expected_user.displayName = "John Doe"

        self.mock_ometa.create_or_update.return_value = expected_user

        # Act
        result = Users.create(create_request)

        # Assert
        self.assertEqual(str(result.id), self.user_id)
        self.assertEqual(result.name, "john.doe")
        self.assertEqual(result.email, "john.doe@company.com")
        self.mock_ometa.create_or_update.assert_called_once_with(create_request)

    def test_retrieve_user_by_id(self):
        """Test retrieving a user by ID"""
        # Arrange
        expected_user = MagicMock(spec=UserEntity)
        expected_user.id = UUID(self.user_id)
        expected_user.name = "john.doe"
        expected_user.description = "Senior Data Engineer"

        self.mock_ometa.get_by_id.return_value = expected_user

        # Act
        result = Users.retrieve(self.user_id)

        # Assert
        self.assertEqual(str(result.id), self.user_id)
        self.assertEqual(result.name, "john.doe")
        self.mock_ometa.get_by_id.assert_called_once_with(
            entity=UserEntity, entity_id=self.user_id, fields=None
        )

    def test_retrieve_user_with_teams(self):
        """Test retrieving user with team memberships"""
        # Arrange
        fields = ["teams", "roles", "owns"]

        # Mock teams
        team1 = EntityReference(
            id=UUID("350e8400-e29b-41d4-a716-446655440000"),
            type="team",
            name="data-engineering",
        )
        team2 = EntityReference(
            id=UUID("450e8400-e29b-41d4-a716-446655440000"),
            type="team",
            name="analytics",
        )

        expected_user = MagicMock(spec=UserEntity)
        expected_user.id = UUID(self.user_id)
        expected_user.name = "john.doe"
        expected_user.teams = [team1, team2]

        self.mock_ometa.get_by_id.return_value = expected_user

        # Act
        result = Users.retrieve(self.user_id, fields=fields)

        # Assert
        self.assertIsNotNone(result.teams)
        self.assertEqual(len(result.teams), 2)
        self.assertEqual(result.teams[0].name, "data-engineering")
        self.mock_ometa.get_by_id.assert_called_once_with(
            entity=UserEntity, entity_id=self.user_id, fields=fields
        )

    def test_retrieve_user_by_name(self):
        """Test retrieving a user by name"""
        # Arrange
        expected_user = MagicMock(spec=UserEntity)
        expected_user.id = UUID(self.user_id)
        expected_user.name = "john.doe"
        expected_user.fullyQualifiedName = self.user_fqn

        self.mock_ometa.get_by_name.return_value = expected_user

        # Act
        result = Users.retrieve_by_name(self.user_fqn)

        # Assert
        self.assertEqual(result.fullyQualifiedName, self.user_fqn)
        self.mock_ometa.get_by_name.assert_called_once_with(
            entity=UserEntity, fqn=self.user_fqn, fields=None
        )

    def test_update_user(self):
        """Test updating a user"""
        # Arrange
        user_to_update = MagicMock(spec=UserEntity)
        user_to_update.id = UUID(self.user_id)
        user_to_update.name = "john.doe"
        user_to_update.description = "Principal Data Engineer"

        # Mock the get_by_id to return the current state
        current_entity = MagicMock(spec=type(user_to_update))
        current_entity.id = (
            user_to_update.id if hasattr(user_to_update, "id") else UUID(self.entity_id)
        )
        self.mock_ometa.get_by_id.return_value = current_entity

        # Mock the patch to return the updated entity
        self.mock_ometa.patch.return_value = user_to_update

        # Act
        result = Users.update(user_to_update)

        # Assert
        self.assertEqual(result.description, "Principal Data Engineer")
        # Verify get_by_id was called to fetch current state
        self.mock_ometa.get_by_id.assert_called_once()
        # Verify patch was called with source and destination
        self.mock_ometa.patch.assert_called_once()

    def test_delete_user(self):
        """Test deleting a user"""
        # Act
        Users.delete(self.user_id, recursive=False, hard_delete=True)

        # Assert
        self.mock_ometa.delete.assert_called_once_with(
            entity=UserEntity, entity_id=self.user_id, recursive=False, hard_delete=True
        )

    def test_user_with_authentication(self):
        """Test user with authentication mechanism"""
        # Arrange
        auth_mechanism = MagicMock(spec=AuthenticationMechanism)
        auth_mechanism.authType = "JWT"

        expected_user = MagicMock(spec=UserEntity)
        expected_user.id = UUID(self.user_id)
        expected_user.authenticationMechanism = auth_mechanism

        self.mock_ometa.get_by_id.return_value = expected_user

        # Act
        result = Users.retrieve(self.user_id, fields=["authenticationMechanism"])

        # Assert
        self.assertIsNotNone(result.authenticationMechanism)
        self.assertEqual(result.authenticationMechanism.authType, "JWT")

    def test_user_with_roles(self):
        """Test user with assigned roles"""
        # Arrange
        role1 = EntityReference(
            id=UUID("550e8400-e29b-41d4-a716-446655440000"),
            type="role",
            name="DataEngineer",
        )
        role2 = EntityReference(
            id=UUID("650e8400-e29b-41d4-a716-446655440000"), type="role", name="Admin"
        )

        expected_user = MagicMock(spec=UserEntity)
        expected_user.id = UUID(self.user_id)
        expected_user.roles = [role1, role2]
        expected_user.isAdmin = True

        self.mock_ometa.get_by_id.return_value = expected_user

        # Act
        result = Users.retrieve(self.user_id, fields=["roles"])

        # Assert
        self.assertIsNotNone(result.roles)
        self.assertEqual(len(result.roles), 2)
        self.assertEqual(result.roles[0].name, "DataEngineer")
        self.assertEqual(result.isAdmin, True)

    def test_user_ownership(self):
        """Test user ownership of assets"""
        # Arrange
        owned_table = EntityReference(
            id=UUID("750e8400-e29b-41d4-a716-446655440000"),
            type="table",
            name="analytics.users",
        )
        owned_dashboard = EntityReference(
            id=UUID("850e8400-e29b-41d4-a716-446655440000"),
            type="dashboard",
            name="user-metrics",
        )

        expected_user = MagicMock(spec=UserEntity)
        expected_user.id = UUID(self.user_id)
        expected_user.owns = [owned_table, owned_dashboard]

        self.mock_ometa.get_by_id.return_value = expected_user

        # Act
        result = Users.retrieve(self.user_id, fields=["owns"])

        # Assert
        self.assertIsNotNone(result.owns)
        self.assertEqual(len(result.owns), 2)
        self.assertEqual(result.owns[0].name, "analytics.users")

    def test_user_with_profile(self):
        """Test user with profile information"""
        # Arrange
        profile = MagicMock()
        profile.images = MagicMock()
        profile.images.image = "https://company.com/avatars/john.doe.png"
        profile.images.image24 = "https://company.com/avatars/john.doe_24.png"

        expected_user = MagicMock(spec=UserEntity)
        expected_user.id = UUID(self.user_id)
        expected_user.profile = profile

        self.mock_ometa.get_by_id.return_value = expected_user

        # Act
        result = Users.retrieve(self.user_id, fields=["profile"])

        # Assert
        self.assertIsNotNone(result.profile)
        self.assertEqual(
            result.profile.images.image, "https://company.com/avatars/john.doe.png"
        )

    def test_list_users(self):
        """Test listing users with pagination"""
        # Arrange
        mock_user1 = MagicMock(spec=UserEntity)
        mock_user1.name = "user1"
        mock_user2 = MagicMock(spec=UserEntity)
        mock_user2.name = "user2"
        mock_user3 = MagicMock(spec=UserEntity)
        mock_user3.name = "user3"

        mock_response = MagicMock()
        mock_response.entities = [mock_user1, mock_user2, mock_user3]

        self.mock_ometa.list_entities.return_value = mock_response

        # Act
        result = Users.list(limit=20, fields=["teams", "roles"])

        # Assert
        self.assertEqual(len(result.entities), 3)
        self.assertEqual(result.entities[0].name, "user1")
        self.mock_ometa.list_entities.assert_called_once()

    def test_user_follows(self):
        """Test user following entities"""
        # Arrange
        followed_table = EntityReference(
            id=UUID("950e8400-e29b-41d4-a716-446655440000"),
            type="table",
            name="important.metrics",
        )

        expected_user = MagicMock(spec=UserEntity)
        expected_user.id = UUID(self.user_id)
        expected_user.follows = [followed_table]

        self.mock_ometa.get_by_id.return_value = expected_user

        # Act
        result = Users.retrieve(self.user_id, fields=["follows"])

        # Assert
        self.assertIsNotNone(result.follows)
        self.assertEqual(result.follows[0].name, "important.metrics")

    def test_error_handling_user_not_found(self):
        """Test error handling when user not found"""
        # Arrange
        self.mock_ometa.get_by_id.side_effect = Exception("User not found")

        # Act & Assert
        with self.assertRaises(Exception) as context:
            Users.retrieve("non-existent-id")

        self.assertIn("User not found", str(context.exception))


if __name__ == "__main__":
    unittest.main()