Spaces:
Sleeping
Sleeping
| """ | |
| Rigorous Tests for Drive Service. | |
| Tests cover: | |
| 1. Initialization and credential loading | |
| 2. OAuth authentication | |
| 3. Folder operations (find, create) | |
| 4. Database upload | |
| 5. Database download | |
| 6. Error handling | |
| Mocks Google API - no real Drive calls. | |
| """ | |
| import pytest | |
| import os | |
| import tempfile | |
| from unittest.mock import patch, MagicMock, PropertyMock | |
| from googleapiclient.errors import HttpError | |
| # ============================================================================= | |
| # 1. Initialization Tests | |
| # ============================================================================= | |
| class TestDriveServiceInit: | |
| """Test DriveService initialization.""" | |
| def test_load_server_credentials(self): | |
| """Load credentials from SERVER_* env vars.""" | |
| with patch.dict(os.environ, { | |
| "SERVER_GOOGLE_CLIENT_ID": "server-client-id", | |
| "SERVER_GOOGLE_CLIENT_SECRET": "server-secret", | |
| "SERVER_GOOGLE_REFRESH_TOKEN": "server-refresh" | |
| }): | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| assert service.client_id == "server-client-id" | |
| assert service.client_secret == "server-secret" | |
| assert service.refresh_token == "server-refresh" | |
| def test_fallback_to_google_credentials(self): | |
| """Fallback to GOOGLE_* env vars when SERVER_* missing.""" | |
| with patch.dict(os.environ, { | |
| "GOOGLE_CLIENT_ID": "google-client-id", | |
| "GOOGLE_CLIENT_SECRET": "google-secret", | |
| "GOOGLE_REFRESH_TOKEN": "google-refresh" | |
| }, clear=True): | |
| # Clear SERVER_* vars | |
| os.environ.pop("SERVER_GOOGLE_CLIENT_ID", None) | |
| os.environ.pop("SERVER_GOOGLE_CLIENT_SECRET", None) | |
| os.environ.pop("SERVER_GOOGLE_REFRESH_TOKEN", None) | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| assert service.client_id == "google-client-id" | |
| def test_init_with_missing_credentials(self): | |
| """Initialize with missing credentials (None values).""" | |
| with patch.dict(os.environ, {}, clear=True): | |
| os.environ.pop("SERVER_GOOGLE_CLIENT_ID", None) | |
| os.environ.pop("GOOGLE_CLIENT_ID", None) | |
| os.environ.pop("SERVER_GOOGLE_CLIENT_SECRET", None) | |
| os.environ.pop("GOOGLE_CLIENT_SECRET", None) | |
| os.environ.pop("SERVER_GOOGLE_REFRESH_TOKEN", None) | |
| os.environ.pop("GOOGLE_REFRESH_TOKEN", None) | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| # Should still initialize, but with None values | |
| assert service.creds is None | |
| assert service.service is None | |
| # ============================================================================= | |
| # 2. Authentication Tests | |
| # ============================================================================= | |
| class TestAuthentication: | |
| """Test authenticate method.""" | |
| def test_authenticate_success(self): | |
| """Successful authentication with valid refresh token.""" | |
| with patch('services.drive_service.Credentials') as mock_creds: | |
| with patch('services.drive_service.build') as mock_build: | |
| mock_cred_instance = MagicMock() | |
| mock_cred_instance.expired = False | |
| mock_creds.return_value = mock_cred_instance | |
| mock_build.return_value = MagicMock() | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.client_id = "test-id" | |
| service.client_secret = "test-secret" | |
| service.refresh_token = "test-token" | |
| result = service.authenticate() | |
| assert result == True | |
| assert service.creds is not None | |
| assert service.service is not None | |
| def test_authenticate_returns_false_when_missing_credentials(self): | |
| """Return False when credentials are missing.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.client_id = None | |
| service.client_secret = "secret" | |
| service.refresh_token = "token" | |
| result = service.authenticate() | |
| assert result == False | |
| def test_authenticate_handles_exception(self): | |
| """Handle authentication exception.""" | |
| with patch('services.drive_service.Credentials') as mock_creds: | |
| mock_creds.side_effect = Exception("Auth failed") | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.client_id = "test-id" | |
| service.client_secret = "test-secret" | |
| service.refresh_token = "test-token" | |
| result = service.authenticate() | |
| assert result == False | |
| def test_authenticate_refreshes_expired_token(self): | |
| """Refresh expired token when needed.""" | |
| with patch('services.drive_service.Credentials') as mock_creds: | |
| with patch('services.drive_service.build') as mock_build: | |
| with patch('services.drive_service.Request') as mock_request: | |
| mock_cred_instance = MagicMock() | |
| mock_cred_instance.expired = True | |
| mock_cred_instance.refresh_token = "has-refresh" | |
| mock_creds.return_value = mock_cred_instance | |
| mock_build.return_value = MagicMock() | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.client_id = "test-id" | |
| service.client_secret = "test-secret" | |
| service.refresh_token = "test-token" | |
| result = service.authenticate() | |
| assert result == True | |
| mock_cred_instance.refresh.assert_called_once() | |
| # ============================================================================= | |
| # 3. Folder Operations Tests | |
| # ============================================================================= | |
| class TestFolderOperations: | |
| """Test folder find/create operations.""" | |
| def test_find_folder_returns_id_when_found(self): | |
| """Find existing folder returns its ID.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.service = MagicMock() | |
| service.service.files.return_value.list.return_value.execute.return_value = { | |
| 'files': [{'id': 'folder-123', 'name': 'apigateway'}] | |
| } | |
| result = service._find_folder() | |
| assert result == 'folder-123' | |
| def test_find_folder_returns_none_when_not_found(self): | |
| """Return None when folder not found.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.service = MagicMock() | |
| service.service.files.return_value.list.return_value.execute.return_value = { | |
| 'files': [] | |
| } | |
| result = service._find_folder() | |
| assert result is None | |
| def test_find_folder_handles_http_error(self): | |
| """Handle HTTP error in folder search.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.service = MagicMock() | |
| # Create mock HttpError | |
| mock_resp = MagicMock() | |
| mock_resp.status = 500 | |
| service.service.files.return_value.list.return_value.execute.side_effect = HttpError( | |
| mock_resp, b"Server error" | |
| ) | |
| result = service._find_folder() | |
| assert result is None | |
| def test_create_folder_returns_id(self): | |
| """Create folder returns new folder ID.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.service = MagicMock() | |
| service.service.files.return_value.create.return_value.execute.return_value = { | |
| 'id': 'new-folder-456' | |
| } | |
| result = service._create_folder() | |
| assert result == 'new-folder-456' | |
| def test_create_folder_handles_error(self): | |
| """Handle error in folder creation.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.service = MagicMock() | |
| mock_resp = MagicMock() | |
| mock_resp.status = 403 | |
| service.service.files.return_value.create.return_value.execute.side_effect = HttpError( | |
| mock_resp, b"Permission denied" | |
| ) | |
| result = service._create_folder() | |
| assert result is None | |
| def test_get_folder_id_creates_if_not_found(self): | |
| """Get folder creates if not found.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service._find_folder = MagicMock(return_value=None) | |
| service._create_folder = MagicMock(return_value='created-folder-789') | |
| result = service._get_folder_id() | |
| assert result == 'created-folder-789' | |
| service._create_folder.assert_called_once() | |
| def test_get_folder_id_returns_existing(self): | |
| """Get folder returns existing folder ID.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service._find_folder = MagicMock(return_value='existing-folder-111') | |
| service._create_folder = MagicMock() | |
| result = service._get_folder_id() | |
| assert result == 'existing-folder-111' | |
| service._create_folder.assert_not_called() | |
| # ============================================================================= | |
| # 4. Upload Database Tests | |
| # ============================================================================= | |
| class TestUploadDb: | |
| """Test upload_db method.""" | |
| def test_upload_new_file(self): | |
| """Upload creates new file when not exists.""" | |
| from services.drive_service import DriveService | |
| with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as f: | |
| f.write(b"test db content") | |
| temp_path = f.name | |
| try: | |
| service = DriveService() | |
| service.DB_FILENAME = temp_path | |
| service.service = MagicMock() | |
| service._get_folder_id = MagicMock(return_value='folder-id') | |
| # File doesn't exist in Drive | |
| service.service.files.return_value.list.return_value.execute.return_value = { | |
| 'files': [] | |
| } | |
| service.service.files.return_value.create.return_value.execute.return_value = { | |
| 'id': 'new-file-id' | |
| } | |
| with patch('services.drive_service.MediaFileUpload') as mock_upload: | |
| mock_upload.return_value = MagicMock() | |
| result = service.upload_db() | |
| assert result == True | |
| service.service.files.return_value.create.assert_called() | |
| finally: | |
| os.unlink(temp_path) | |
| def test_upload_updates_existing_file(self): | |
| """Upload updates existing file.""" | |
| from services.drive_service import DriveService | |
| with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as f: | |
| f.write(b"updated db content") | |
| temp_path = f.name | |
| try: | |
| service = DriveService() | |
| service.DB_FILENAME = temp_path | |
| service.service = MagicMock() | |
| service._get_folder_id = MagicMock(return_value='folder-id') | |
| # File exists in Drive | |
| service.service.files.return_value.list.return_value.execute.return_value = { | |
| 'files': [{'id': 'existing-file-id'}] | |
| } | |
| with patch('services.drive_service.MediaFileUpload') as mock_upload: | |
| mock_upload.return_value = MagicMock() | |
| result = service.upload_db() | |
| assert result == True | |
| service.service.files.return_value.update.assert_called() | |
| finally: | |
| os.unlink(temp_path) | |
| def test_upload_returns_false_if_db_missing(self): | |
| """Return False if database file doesn't exist.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.DB_FILENAME = '/nonexistent/path/db.sqlite' | |
| service.service = MagicMock() | |
| result = service.upload_db() | |
| assert result == False | |
| def test_upload_returns_false_if_folder_fails(self): | |
| """Return False if folder creation fails.""" | |
| from services.drive_service import DriveService | |
| with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as f: | |
| f.write(b"test") | |
| temp_path = f.name | |
| try: | |
| service = DriveService() | |
| service.DB_FILENAME = temp_path | |
| service.service = MagicMock() | |
| service._get_folder_id = MagicMock(return_value=None) | |
| result = service.upload_db() | |
| assert result == False | |
| finally: | |
| os.unlink(temp_path) | |
| def test_upload_handles_http_error(self): | |
| """Handle HTTP error during upload.""" | |
| from services.drive_service import DriveService | |
| with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as f: | |
| f.write(b"test") | |
| temp_path = f.name | |
| try: | |
| service = DriveService() | |
| service.DB_FILENAME = temp_path | |
| service.service = MagicMock() | |
| service._get_folder_id = MagicMock(return_value='folder-id') | |
| service.service.files.return_value.list.return_value.execute.return_value = { | |
| 'files': [] | |
| } | |
| mock_resp = MagicMock() | |
| mock_resp.status = 500 | |
| service.service.files.return_value.create.return_value.execute.side_effect = HttpError( | |
| mock_resp, b"Server error" | |
| ) | |
| with patch('services.drive_service.MediaFileUpload'): | |
| result = service.upload_db() | |
| assert result == False | |
| finally: | |
| os.unlink(temp_path) | |
| def test_upload_auto_authenticates(self): | |
| """Auto-authenticate if not authenticated.""" | |
| from services.drive_service import DriveService | |
| with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as f: | |
| f.write(b"test") | |
| temp_path = f.name | |
| try: | |
| service = DriveService() | |
| service.DB_FILENAME = temp_path | |
| service.service = None | |
| service.authenticate = MagicMock(return_value=False) | |
| result = service.upload_db() | |
| assert result == False | |
| service.authenticate.assert_called_once() | |
| finally: | |
| os.unlink(temp_path) | |
| # ============================================================================= | |
| # 5. Download Database Tests | |
| # ============================================================================= | |
| class TestDownloadDb: | |
| """Test download_db method.""" | |
| def test_download_existing_file(self): | |
| """Download existing database file.""" | |
| from services.drive_service import DriveService | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| db_path = os.path.join(temp_dir, 'test.db') | |
| service = DriveService() | |
| service.DB_FILENAME = db_path | |
| service.service = MagicMock() | |
| service._find_folder = MagicMock(return_value='folder-id') | |
| service.service.files.return_value.list.return_value.execute.return_value = { | |
| 'files': [{'id': 'file-id'}] | |
| } | |
| # Mock downloader | |
| with patch('services.drive_service.MediaIoBaseDownload') as mock_downloader: | |
| mock_instance = MagicMock() | |
| mock_instance.next_chunk.return_value = (MagicMock(), True) # Done immediately | |
| mock_downloader.return_value = mock_instance | |
| with patch('io.FileIO'): | |
| result = service.download_db() | |
| assert result == True | |
| def test_download_returns_false_if_folder_not_found(self): | |
| """Return False if folder not found.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.service = MagicMock() | |
| service._find_folder = MagicMock(return_value=None) | |
| result = service.download_db() | |
| assert result == False | |
| def test_download_returns_false_if_file_not_found(self): | |
| """Return False if database file not found in Drive.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.service = MagicMock() | |
| service._find_folder = MagicMock(return_value='folder-id') | |
| service.service.files.return_value.list.return_value.execute.return_value = { | |
| 'files': [] | |
| } | |
| result = service.download_db() | |
| assert result == False | |
| def test_download_handles_http_error(self): | |
| """Handle HTTP error during download.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.service = MagicMock() | |
| service._find_folder = MagicMock(return_value='folder-id') | |
| service.service.files.return_value.list.return_value.execute.return_value = { | |
| 'files': [{'id': 'file-id'}] | |
| } | |
| mock_resp = MagicMock() | |
| mock_resp.status = 500 | |
| service.service.files.return_value.get_media.side_effect = HttpError( | |
| mock_resp, b"Server error" | |
| ) | |
| result = service.download_db() | |
| assert result == False | |
| def test_download_auto_authenticates(self): | |
| """Auto-authenticate if not authenticated.""" | |
| from services.drive_service import DriveService | |
| service = DriveService() | |
| service.service = None | |
| service.authenticate = MagicMock(return_value=False) | |
| result = service.download_db() | |
| assert result == False | |
| service.authenticate.assert_called_once() | |
| # ============================================================================= | |
| # 6. Integration-like Tests | |
| # ============================================================================= | |
| class TestDriveServiceIntegration: | |
| """Test complete flows with mocked API.""" | |
| def test_full_upload_flow(self): | |
| """Test complete upload flow from auth to upload.""" | |
| from services.drive_service import DriveService | |
| with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as f: | |
| f.write(b"test db content") | |
| temp_path = f.name | |
| try: | |
| with patch('services.drive_service.Credentials') as mock_creds: | |
| with patch('services.drive_service.build') as mock_build: | |
| mock_cred_instance = MagicMock() | |
| mock_cred_instance.expired = False | |
| mock_creds.return_value = mock_cred_instance | |
| mock_service = MagicMock() | |
| mock_build.return_value = mock_service | |
| # Folder exists | |
| mock_service.files.return_value.list.return_value.execute.side_effect = [ | |
| {'files': [{'id': 'folder-id', 'name': 'apigateway'}]}, # find folder | |
| {'files': []} # file not in folder | |
| ] | |
| mock_service.files.return_value.create.return_value.execute.return_value = { | |
| 'id': 'new-file-id' | |
| } | |
| service = DriveService() | |
| service.DB_FILENAME = temp_path | |
| service.client_id = "test" | |
| service.client_secret = "test" | |
| service.refresh_token = "test" | |
| with patch('services.drive_service.MediaFileUpload'): | |
| result = service.upload_db() | |
| assert result == True | |
| finally: | |
| os.unlink(temp_path) | |
| if __name__ == "__main__": | |
| pytest.main([__file__, "-v"]) | |