tests/test_federation.py

import pytest
from datetime import datetime

from blossomtune_gradio import federation as fed
from blossomtune_gradio.database import Request


@pytest.fixture
def mock_mail(mocker):
    """Fixture to mock the mail module."""
    return mocker.patch("blossomtune_gradio.mail.send_activation_email")


def test_generate_participant_id():
    """Test the generation of a participant ID."""
    pid = fed.generate_participant_id()
    assert isinstance(pid, str)
    assert len(pid) == 6
    assert pid.isalnum() and pid.isupper()


def test_generate_activation_code():
    """Test the generation of an activation code."""
    code = fed.generate_activation_code()
    assert isinstance(code, str)
    assert len(code) == 8
    assert code.isalnum() and code.isupper()


class TestCheckParticipantStatus:
    """Test suite for the check_participant_status function using SQLAlchemy."""

    def test_new_user_registration_success(self, db_session, mock_settings, mock_mail):
        """Verify successful registration for a new user."""
        mock_mail.return_value = (True, "")
        approved, message, download = fed.check_participant_status(
            "new_user", "hello@ethicalabs.ai", ""
        )
        assert approved is False
        assert download is None
        assert message == "mock_registration_submitted_md"

        # Verify the user was added to the database
        request = db_session.query(Request).filter_by(hf_handle="new_user").first()
        assert request is not None
        assert request.email == "hello@ethicalabs.ai"

    def test_new_user_invalid_email(self, db_session, mock_settings):
        """Verify registration fails with an invalid email."""
        approved, message, download = fed.check_participant_status(
            "user", "invalid-email", ""
        )
        assert approved is False
        assert download is None
        assert message == "mock_invalid_email_md"

    def test_new_user_federation_full(self, db_session, mock_settings, mocker):
        """Verify registration fails when the federation is full."""
        mocker.patch("blossomtune_gradio.config.MAX_NUM_NODES", 0)
        approved, message, download = fed.check_participant_status(
            "another_user", "another@example.com", ""
        )
        assert approved is False
        assert download is None
        assert message == "mock_federation_full_md"

    def test_user_activation_success(self, db_session, mock_settings):
        """Verify a user can successfully activate their account."""
        # Setup: Add a pending, non-activated user
        pending_user = Request(
            participant_id="PID123",
            hf_handle="test_user",
            email="test@example.com",
            activation_code="ABCDEF",
            is_activated=0,
        )
        db_session.add(pending_user)
        db_session.commit()

        approved, message, download = fed.check_participant_status(
            "test_user", "test@example.com", "ABCDEF"
        )
        assert approved is False
        assert download is None
        assert message == "mock_activation_successful_md"
        # Verify the user is now activated
        activated_user = (
            db_session.query(Request).filter_by(hf_handle="test_user").first()
        )
        assert activated_user.is_activated == 1

    def test_user_activation_invalid_code(self, db_session, mock_settings):
        """Verify activation fails with an invalid code."""
        pending_user = Request(
            participant_id="PID123",
            hf_handle="test_user",
            email="test@example.com",
            activation_code="ABCDEF",
            is_activated=0,
        )
        db_session.add(pending_user)
        db_session.commit()

        approved, message, download = fed.check_participant_status(
            "test_user", "test@example.com", "WRONGCODE"
        )
        assert approved is False
        assert download is None
        assert message == "mock_activation_invalid_md"

    def test_status_check_approved_unmanaged(self, db_session, mock_settings):
        """Verify the status check for an approved user."""
        approved_user = Request(
            participant_id="PID456",
            status="approved",
            hf_handle="approved_user",
            email="approved@example.com",
            activation_code="GHIJKL",
            is_activated=1,
            partition_id=5,
        )
        db_session.add(approved_user)
        db_session.commit()

        approved, message, download = fed.check_participant_status(
            "approved_user", "approved@example.com", "GHIJKL"
        )
        assert approved is False
        assert download is None  # fed.manage_request is not used. download is None.
        assert "An error occurred" in message


class TestManageRequest:
    """Test suite for the manage_request function using SQLAlchemy."""

    def test_approve_success(self, db_session):
        """Verify successful approval of a participant."""
        pending_user = Request(
            participant_id="PENDING1",
            status="pending",
            hf_handle="pending_user",
            email="pending@example.com",
            is_activated=1,
        )
        db_session.add(pending_user)
        db_session.commit()

        success, message = fed.manage_request("PENDING1", "10", "approve")
        assert success is True
        assert (
            "Participant PENDING1 approved. Keys generated and registry updated."
            in message
        )

        # Verify status in DB
        updated_user = (
            db_session.query(Request).filter_by(participant_id="PENDING1").first()
        )
        assert updated_user.status == "approved"
        assert updated_user.partition_id == 10

    def test_approve_not_activated(self, db_session, mock_settings):
        """Verify approval fails if the user is not activated."""
        pending_user = Request(
            participant_id="PENDING2",
            status="pending",
            hf_handle="pending_user2",
            is_activated=0,
        )
        db_session.add(pending_user)
        db_session.commit()
        success, message = fed.manage_request("PENDING2", "11", "approve")
        assert success is False
        assert message == "mock_participant_not_activated_warning_md"

    def test_deny_success(self, db_session):
        """Verify successful denial of a participant."""
        pending_user = Request(
            participant_id="PENDING3",
            status="pending",
            hf_handle="pending_user3",
            is_activated=1,
        )
        db_session.add(pending_user)
        db_session.commit()
        success, message = fed.manage_request("PENDING3", "", "deny")
        assert success is True
        assert "Participant PENDING3 denied. Their access has been revoked." in message

        # Verify status in DB
        updated_user = (
            db_session.query(Request).filter_by(participant_id="PENDING3").first()
        )
        assert updated_user.status == "denied"


def test_get_next_partition_id(db_session):
    """Verify the logic for finding the next available partition ID."""
    # No approved users yet
    assert fed.get_next_partion_id() == 0

    # Add some approved users with assigned partitions
    db_session.add(
        Request(
            participant_id="P1",
            status="approved",
            partition_id=0,
            timestamp=datetime.utcnow(),
        )
    )
    db_session.add(
        Request(
            participant_id="P2",
            status="approved",
            partition_id=1,
            timestamp=datetime.utcnow(),
        )
    )
    db_session.commit()
    assert fed.get_next_partion_id() == 2

    # Add another user, skipping a partition ID
    db_session.add(
        Request(
            participant_id="P3",
            status="approved",
            partition_id=3,
            timestamp=datetime.utcnow(),
        )
    )
    db_session.commit()
    assert fed.get_next_partion_id() == 2