dify/api/tests/unit_tests/models/test_comment_models.py

from uuid import uuid4

import pytest
from sqlalchemy.orm import Session

from models.account import Account
from models.comment import WorkflowComment, WorkflowCommentMention, WorkflowCommentReply


def _account(name: str) -> Account:
    account = Account(name=name, email=f"{name.lower()}@example.com")
    account.id = str(uuid4())
    return account


def _comment(created_by: str, resolved_by: str | None = None) -> WorkflowComment:
    return WorkflowComment(
        created_by=created_by,
        resolved_by=resolved_by,
        content="hello",
        position_x=1,
        position_y=2,
        tenant_id=str(uuid4()),
        app_id=str(uuid4()),
    )


COMMENT_MODELS = (Account, WorkflowComment, WorkflowCommentReply, WorkflowCommentMention)


@pytest.mark.parametrize("sqlite_session", [COMMENT_MODELS], indirect=True)
def test_workflow_comment_account_accessors_and_cache(sqlite_session: Session) -> None:
    created_account = _account("Creator")
    resolved_account = _account("Resolver")
    comment = _comment(created_account.id, resolved_account.id)
    sqlite_session.add_all([created_account, resolved_account, comment])
    sqlite_session.commit()

    assert comment.created_by_account(sqlite_session) is created_account
    assert comment.resolved_by_account(sqlite_session) is resolved_account

    comment.cache_created_by_account(created_account)
    comment.cache_resolved_by_account(resolved_account)
    sqlite_session.delete(created_account)
    sqlite_session.delete(resolved_account)
    sqlite_session.commit()

    assert comment.created_by_account(sqlite_session) is created_account
    assert comment.resolved_by_account(sqlite_session) is resolved_account

    comment_without_resolver = _comment(str(uuid4()))
    sqlite_session.add(comment_without_resolver)
    sqlite_session.commit()
    assert comment_without_resolver.resolved_by_account(sqlite_session) is None


@pytest.mark.parametrize("sqlite_session", [COMMENT_MODELS], indirect=True)
def test_workflow_comment_counts_and_participants(sqlite_session: Session) -> None:
    account_1 = _account("Creator")
    account_2 = _account("Replier")
    account_3 = _account("Mentioned")
    missing_account_id = str(uuid4())
    comment = _comment(account_1.id)
    comment.replies = [
        WorkflowCommentReply(comment_id=comment.id, content="reply-1", created_by=account_2.id),
        WorkflowCommentReply(comment_id=comment.id, content="reply-2", created_by=account_2.id),
    ]
    comment.mentions = [
        WorkflowCommentMention(comment_id=comment.id, mentioned_user_id=account_3.id),
        WorkflowCommentMention(comment_id=comment.id, mentioned_user_id=missing_account_id),
    ]
    sqlite_session.add_all([account_1, account_2, account_3, comment])
    sqlite_session.commit()

    participants = comment.participants(sqlite_session)

    assert comment.reply_count == 2
    assert comment.mention_count == 2
    assert {account.id for account in participants} == {account_1.id, account_2.id, account_3.id}


@pytest.mark.parametrize("sqlite_session", [COMMENT_MODELS], indirect=True)
def test_workflow_comment_participants_use_cached_accounts(sqlite_session: Session) -> None:
    account_1 = _account("Creator")
    account_2 = _account("Replier")
    account_3 = _account("Mentioned")
    comment = _comment(account_1.id)
    reply = WorkflowCommentReply(comment_id=comment.id, content="reply-1", created_by=account_2.id)
    mention = WorkflowCommentMention(comment_id=comment.id, mentioned_user_id=account_3.id)
    comment.replies = [reply]
    comment.mentions = [mention]
    sqlite_session.add_all([account_1, account_2, account_3, comment])
    sqlite_session.commit()

    comment.cache_created_by_account(account_1)
    reply.cache_created_by_account(account_2)
    mention.cache_mentioned_user_account(account_3)
    sqlite_session.delete(account_1)
    sqlite_session.delete(account_2)
    sqlite_session.delete(account_3)
    sqlite_session.commit()

    assert {account.id for account in comment.participants(sqlite_session)} == {
        account_1.id,
        account_2.id,
        account_3.id,
    }


@pytest.mark.parametrize("sqlite_session", [COMMENT_MODELS], indirect=True)
def test_reply_and_mention_account_accessors_and_cache(sqlite_session: Session) -> None:
    reply_account = _account("Replier")
    mention_account = _account("Mentioned")
    comment = _comment(reply_account.id)
    reply = WorkflowCommentReply(comment_id=comment.id, content="reply", created_by=reply_account.id)
    mention = WorkflowCommentMention(comment_id=comment.id, mentioned_user_id=mention_account.id)
    comment.replies = [reply]
    comment.mentions = [mention]
    sqlite_session.add_all([reply_account, mention_account, comment])
    sqlite_session.commit()

    assert reply.created_by_account(sqlite_session) is reply_account
    assert mention.mentioned_user_account(sqlite_session) is mention_account

    reply.cache_created_by_account(reply_account)
    mention.cache_mentioned_user_account(mention_account)
    sqlite_session.delete(reply_account)
    sqlite_session.delete(mention_account)
    sqlite_session.commit()

    assert reply.created_by_account(sqlite_session) is reply_account
    assert mention.mentioned_user_account(sqlite_session) is mention_account