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()