dify/api/tests/unit_tests/services/test_account_education_service.py

from __future__ import annotations

from datetime import UTC, datetime
from unittest.mock import Mock

import pytest

from machinery.context import RequestContext
from services.account_education_service import (
    AccountEducationGateway,
    AccountEducationRateLimiter,
    AccountEducationService,
)
from services.account_errors import EducationRateLimitExceededError
from services.account_ports import AccountRepository
from services.entities.account_entities import (
    AccountEducationActivation,
    AccountEducationAutocomplete,
    AccountEducationStatus,
    AccountEducationVerification,
    AccountSnapshot,
)


def _context() -> RequestContext:
    return RequestContext(
        request_id="request-1",
        trace_id="trace-1",
        account_id="account-1",
        active_workspace_id="workspace-1",
    )


def _account() -> AccountSnapshot:
    return AccountSnapshot(
        id="account-1",
        name="Student",
        email="student@example.edu",
        avatar=None,
        is_password_set=True,
        interface_language="en-US",
        interface_theme="light",
        timezone="UTC",
        last_login_at=None,
        last_login_ip=None,
        status="active",
        initialized_at=datetime(2026, 1, 1),
        created_at=datetime(2026, 1, 1),
    )


def _rate_limiter() -> Mock:
    rate_limiter = Mock(spec=AccountEducationRateLimiter)
    rate_limiter.is_rate_limited.return_value = False
    return rate_limiter


def test_verify_reads_account_before_billing_gateway_call() -> None:
    accounts = Mock(spec=AccountRepository)
    accounts.get.return_value = _account()
    education = Mock(spec=AccountEducationGateway)
    verification_rate_limiter = _rate_limiter()
    activation_rate_limiter = _rate_limiter()
    events: list[str] = []
    verification_rate_limiter.is_rate_limited.side_effect = lambda _key: events.append("check") or False
    verification_rate_limiter.increment_rate_limit.side_effect = lambda _key: events.append("increment")
    education.verify.side_effect = lambda **_kwargs: (
        events.append("verify") or AccountEducationVerification(token="education-token")
    )
    service = AccountEducationService(
        accounts=accounts,
        education=education,
        verification_rate_limiter=verification_rate_limiter,
        activation_rate_limiter=activation_rate_limiter,
    )

    result = service.verify(_context())

    assert result == AccountEducationVerification(token="education-token")
    assert events == ["check", "increment", "verify"]
    accounts.get.assert_called_once_with("account-1")
    verification_rate_limiter.is_rate_limited.assert_called_once_with("student@example.edu")
    verification_rate_limiter.increment_rate_limit.assert_called_once_with("student@example.edu")
    education.verify.assert_called_once_with(account_id="account-1")
    activation_rate_limiter.is_rate_limited.assert_not_called()


def test_status_and_autocomplete_delegate_framework_neutral_contracts() -> None:
    accounts = Mock(spec=AccountRepository)
    education = Mock(spec=AccountEducationGateway)
    status = AccountEducationStatus(
        result=True,
        is_student=True,
        expire_at=datetime(2027, 1, 1, tzinfo=UTC),
        allow_refresh=False,
    )
    autocomplete = AccountEducationAutocomplete(data=("Example University",), curr_page=0, has_next=False)
    education.status.return_value = status
    education.autocomplete.return_value = autocomplete
    service = AccountEducationService(
        accounts=accounts,
        education=education,
        verification_rate_limiter=_rate_limiter(),
        activation_rate_limiter=_rate_limiter(),
    )

    assert service.status(_context()) == status
    assert service.autocomplete(_context(), keywords="Example", page=0, limit=20) == autocomplete
    education.status.assert_called_once_with("account-1")
    education.autocomplete.assert_called_once_with(keywords="Example", page=0, limit=20)


def test_activate_delegates_account_and_workspace_context() -> None:
    accounts = Mock(spec=AccountRepository)
    accounts.get.return_value = _account()
    education = Mock(spec=AccountEducationGateway)
    activation = AccountEducationActivation(message="success")
    verification_rate_limiter = _rate_limiter()
    activation_rate_limiter = _rate_limiter()
    events: list[str] = []
    activation_rate_limiter.is_rate_limited.side_effect = lambda _key: events.append("check") or False
    activation_rate_limiter.increment_rate_limit.side_effect = lambda _key: events.append("increment")
    education.activate.side_effect = lambda **_kwargs: events.append("activate") or activation
    service = AccountEducationService(
        accounts=accounts,
        education=education,
        verification_rate_limiter=verification_rate_limiter,
        activation_rate_limiter=activation_rate_limiter,
    )

    result = service.activate(
        _context(),
        token="education-token",
        institution="Dify University",
        role="Student",
    )

    assert result == activation
    assert events == ["check", "increment", "activate"]
    activation_rate_limiter.is_rate_limited.assert_called_once_with("student@example.edu")
    activation_rate_limiter.increment_rate_limit.assert_called_once_with("student@example.edu")
    verification_rate_limiter.is_rate_limited.assert_not_called()
    education.activate.assert_called_once_with(
        account_id="account-1",
        tenant_id="workspace-1",
        token="education-token",
        institution="Dify University",
        role="Student",
    )


def test_verify_rejects_rate_limited_request() -> None:
    accounts = Mock(spec=AccountRepository)
    accounts.get.return_value = _account()
    education = Mock(spec=AccountEducationGateway)
    verification_rate_limiter = _rate_limiter()
    verification_rate_limiter.is_rate_limited.return_value = True
    service = AccountEducationService(
        accounts=accounts,
        education=education,
        verification_rate_limiter=verification_rate_limiter,
        activation_rate_limiter=_rate_limiter(),
    )

    with pytest.raises(EducationRateLimitExceededError):
        service.verify(_context())

    verification_rate_limiter.increment_rate_limit.assert_not_called()
    education.verify.assert_not_called()


def test_activate_rejects_rate_limited_request() -> None:
    accounts = Mock(spec=AccountRepository)
    accounts.get.return_value = _account()
    education = Mock(spec=AccountEducationGateway)
    activation_rate_limiter = _rate_limiter()
    activation_rate_limiter.is_rate_limited.return_value = True
    service = AccountEducationService(
        accounts=accounts,
        education=education,
        verification_rate_limiter=_rate_limiter(),
        activation_rate_limiter=activation_rate_limiter,
    )

    with pytest.raises(EducationRateLimitExceededError):
        service.activate(
            _context(),
            token="education-token",
            institution="Dify University",
            role="Student",
        )

    activation_rate_limiter.increment_rate_limit.assert_not_called()
    education.activate.assert_not_called()