Spaces:
Sleeping
Sleeping
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
|