File size: 9,214 Bytes
e60e679
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5f72bf5
e60e679
 
 
5f72bf5
e60e679
 
 
 
 
 
 
 
 
5f72bf5
 
 
e60e679
5f72bf5
e60e679
 
 
 
 
5f72bf5
e60e679
 
 
5f72bf5
e60e679
 
 
 
 
 
 
 
 
 
 
 
5f72bf5
e60e679
 
 
5f72bf5
e60e679
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5f72bf5
e60e679
 
 
5f72bf5
e60e679
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5f72bf5
e60e679
 
 
5f72bf5
e60e679
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
import pytest
from unittest.mock import MagicMock

from blossomtune_gradio import federation as fed


@pytest.fixture
def mock_settings(mocker):
    """Fixture to mock the settings module."""
    # The lambda returns a formatted string to simulate Jinja2's behavior
    mock_get = MagicMock(
        side_effect=lambda key, **kwargs: f"mock_{key}".format(**kwargs)
    )
    mocker.patch("blossomtune_gradio.federation.settings.get_text", mock_get)
    return mock_get


@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."""

    def test_new_user_registration_success(
        self, in_memory_db, mock_settings, mock_mail
    ):
        """Verify successful registration for a new user."""
        mock_mail.return_value = (True, "")
        success, message, download = fed.check_participant_status(
            "new_user", "new@example.com", ""
        )
        assert success is True
        assert download is None
        assert message == "mock_registration_submitted_md"

        # Verify the user was added to the database
        cursor = in_memory_db.cursor()
        cursor.execute("SELECT hf_handle FROM requests WHERE hf_handle = 'new_user'")
        assert cursor.fetchone() is not None

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

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

    def test_user_activation_success(self, in_memory_db, mock_settings):
        """Verify a user can successfully activate their account."""
        # Setup: Add a pending, non-activated user
        cursor = in_memory_db.cursor()
        cursor.execute(
            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
            ("PID123", "pending", "now", "test_user", "test@example.com", "ABCDEF", 0),
        )
        in_memory_db.commit()

        success, message, download = fed.check_participant_status(
            "test_user", "test@example.com", "ABCDEF"
        )
        assert success is True
        assert download is None
        assert message == "mock_activation_successful_md"
        # Verify the user is now activated
        cursor.execute(
            "SELECT is_activated FROM requests WHERE hf_handle = 'test_user'"
        )
        assert cursor.fetchone()[0] == 1

    def test_user_activation_invalid_code(self, in_memory_db, mock_settings):
        """Verify activation fails with an invalid code."""
        # Setup: Add a pending, non-activated user
        cursor = in_memory_db.cursor()
        cursor.execute(
            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
            ("PID123", "pending", "now", "test_user", "test@example.com", "ABCDEF", 0),
        )
        in_memory_db.commit()

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

    def test_status_check_approved(self, in_memory_db, mock_settings):
        """Verify the status check for an approved user."""
        cursor = in_memory_db.cursor()
        cursor.execute(
            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated, partition_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            (
                "PID456",
                "approved",
                "now",
                "approved_user",
                "approved@example.com",
                "GHIJKL",
                1,
                5,
            ),
        )
        in_memory_db.commit()
        success, message, download = fed.check_participant_status(
            "approved_user", "approved@example.com", "GHIJKL"
        )
        assert success is True
        assert download is not None
        assert "mock_status_approved_md" in message


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

    def test_approve_success(self, in_memory_db):
        """Verify successful approval of a participant."""
        # Setup: Add an activated, pending user
        cursor = in_memory_db.cursor()
        cursor.execute(
            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
            (
                "PENDING1",
                "pending",
                "now",
                "pending_user",
                "pending@example.com",
                "CODE",
                1,
            ),
        )
        in_memory_db.commit()

        success, message = fed.manage_request("PENDING1", "10", "approve")
        assert success is True
        assert "is allowed to join" in message

        # Verify status in DB
        cursor.execute(
            "SELECT status, partition_id FROM requests WHERE participant_id = 'PENDING1'"
        )
        status, partition_id = cursor.fetchone()
        assert status == "approved"
        assert partition_id == 10

    def test_approve_not_activated(self, in_memory_db, mock_settings):
        """Verify approval fails if the user is not activated."""
        cursor = in_memory_db.cursor()
        cursor.execute(
            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
            (
                "PENDING2",
                "pending",
                "now",
                "pending_user2",
                "pending2@example.com",
                "CODE",
                0,
            ),
        )
        in_memory_db.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, in_memory_db):
        """Verify successful denial of a participant."""
        cursor = in_memory_db.cursor()
        cursor.execute(
            "INSERT INTO requests (participant_id, status, timestamp, hf_handle, email, activation_code, is_activated) VALUES (?, ?, ?, ?, ?, ?, ?)",
            (
                "PENDING3",
                "pending",
                "now",
                "pending_user3",
                "pending3@example.com",
                "CODE",
                1,
            ),
        )
        in_memory_db.commit()
        success, message = fed.manage_request("PENDING3", "", "deny")
        assert success is True
        assert "is not allowed to join" in message

        # Verify status in DB
        cursor.execute("SELECT status FROM requests WHERE participant_id = 'PENDING3'")
        assert cursor.fetchone()[0] == "denied"


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

    # Add some approved users with assigned partitions, including the required timestamp
    cursor.execute(
        "INSERT INTO requests (participant_id, status, timestamp, partition_id) VALUES (?, ?, ?, ?)",
        ("P1", "approved", "now", 0),
    )
    cursor.execute(
        "INSERT INTO requests (participant_id, status, timestamp, partition_id) VALUES (?, ?, ?, ?)",
        ("P2", "approved", "now", 1),
    )
    in_memory_db.commit()
    assert fed.get_next_partion_id() == 2

    # Add another user, skipping a partition ID
    cursor.execute(
        "INSERT INTO requests (participant_id, status, timestamp, partition_id) VALUES (?, ?, ?, ?)",
        ("P3", "approved", "now", 3),
    )
    in_memory_db.commit()
    assert fed.get_next_partion_id() == 2