Spaces:
Sleeping
Sleeping
File size: 6,342 Bytes
b3e25ad |
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 |
import unittest
import os
import sys
import json
from unittest.mock import patch, MagicMock
# Add parent directory to path for imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# Import modules to test
from auth import HuggingFaceAuth
class MockDB:
"""Mock database for testing."""
def __init__(self):
self.users = {}
def add_user(self, username, hf_user_id, is_admin=False):
user_id = len(self.users) + 1
self.users[hf_user_id] = {
'id': user_id,
'username': username,
'hf_user_id': hf_user_id,
'is_admin': is_admin
}
return user_id
def get_user(self, hf_user_id):
return self.users.get(hf_user_id)
def get_user_by_username(self, username):
for user in self.users.values():
if user['username'] == username:
return user
return None
def can_submit_today(self, user_id):
return True
def update_submission_date(self, user_id):
pass
class MockRequest:
"""Mock Gradio request for testing."""
def __init__(self, headers=None, cookies=None):
self.headers = headers or {}
self.cookies = cookies or {}
class TestHuggingFaceAuth(unittest.TestCase):
"""Test the HuggingFaceAuth class."""
def setUp(self):
"""Set up test environment."""
self.db = MockDB()
# Set environment variables for testing
os.environ['SPACE_ID'] = 'test_space'
os.environ['OAUTH_CLIENT_ID'] = 'test_client_id'
os.environ['OAUTH_CLIENT_SECRET'] = 'test_client_secret'
os.environ['OAUTH_SCOPES'] = 'openid profile'
os.environ['OPENID_PROVIDER_URL'] = 'https://huggingface.co'
self.auth = HuggingFaceAuth(self.db)
def tearDown(self):
"""Clean up after tests."""
# Remove environment variables
for var in ['SPACE_ID', 'OAUTH_CLIENT_ID', 'OAUTH_CLIENT_SECRET', 'OAUTH_SCOPES', 'OPENID_PROVIDER_URL']:
if var in os.environ:
del os.environ[var]
def test_init(self):
"""Test initialization."""
self.assertEqual(self.auth.admin_username, 'Quazim0t0')
self.assertTrue(self.auth.running_in_space)
self.assertEqual(self.auth.client_id, 'test_client_id')
self.assertEqual(self.auth.client_secret, 'test_client_secret')
self.assertEqual(self.auth.oauth_scopes, 'openid profile')
self.assertEqual(self.auth.openid_provider_url, 'https://huggingface.co')
@patch('auth.HfApi')
def test_login_user(self, mock_hf_api):
"""Test login_user method."""
# Mock HfApi.whoami
mock_instance = mock_hf_api.return_value
mock_instance.whoami.return_value = {
'id': 'test_user_id',
'name': 'Test User'
}
# Test successful login
user = self.auth.login_user('test_token')
self.assertIsNotNone(user)
self.assertEqual(user['username'], 'Test User')
self.assertEqual(user['hf_user_id'], 'test_user_id')
self.assertEqual(user['token'], 'test_token')
# Test admin login
mock_instance.whoami.return_value = {
'id': 'admin_user_id',
'name': 'Quazim0t0'
}
user = self.auth.login_user('admin_token')
self.assertIsNotNone(user)
self.assertEqual(user['username'], 'Quazim0t0')
self.assertTrue(user['is_admin'])
# Test failed login
mock_instance.whoami.return_value = None
user = self.auth.login_user('invalid_token')
self.assertIsNone(user)
def test_check_login_space_oauth(self):
"""Test check_login method with Space OAuth."""
# Test with HF-User header
request = MockRequest(headers={'HF-User': 'Test User'})
user = self.auth.check_login(request)
self.assertIsNotNone(user)
self.assertEqual(user['username'], 'Test User')
# Test with admin username
request = MockRequest(headers={'HF-User': 'Quazim0t0'})
user = self.auth.check_login(request)
self.assertIsNotNone(user)
self.assertEqual(user['username'], 'Quazim0t0')
self.assertTrue(user['is_admin'])
@patch('auth.HfApi')
def test_check_login_token(self, mock_hf_api):
"""Test check_login method with token."""
# Remove SPACE_ID to test token-based auth
del os.environ['SPACE_ID']
self.auth.running_in_space = False
# Mock HfApi.whoami
mock_instance = mock_hf_api.return_value
mock_instance.whoami.return_value = {
'id': 'test_user_id',
'name': 'Test User'
}
# Add user to database
self.db.add_user('Test User', 'test_user_id')
# Test with token in cookie
request = MockRequest(cookies={'hf_token': 'test_token'})
user = self.auth.check_login(request)
self.assertIsNotNone(user)
self.assertEqual(user['hf_user_id'], 'test_user_id')
# Test with token in header
request = MockRequest(headers={'HF-Token': 'test_token'})
user = self.auth.check_login(request)
self.assertIsNotNone(user)
self.assertEqual(user['hf_user_id'], 'test_user_id')
# Test with invalid token
mock_instance.whoami.return_value = None
request = MockRequest(cookies={'hf_token': 'invalid_token'})
user = self.auth.check_login(request)
self.assertIsNone(user)
def test_get_oauth_login_url(self):
"""Test get_oauth_login_url method."""
# Test with default redirect URI
url = self.auth.get_oauth_login_url()
self.assertIsNotNone(url)
self.assertIn('client_id=test_client_id', url)
self.assertIn('scope=openid%20profile', url)
self.assertIn('response_type=code', url)
# Test with custom redirect URI
url = self.auth.get_oauth_login_url('https://example.com/callback')
self.assertIsNotNone(url)
self.assertIn('redirect_uri=https://example.com/callback', url)
if __name__ == '__main__':
unittest.main()
|