test_login.py

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