Skip to content

Commit

Permalink
Add test for cancellation failure
Browse files Browse the repository at this point in the history
  • Loading branch information
sunank200 committed Apr 15, 2024
1 parent 1b234c4 commit 5371506
Showing 1 changed file with 29 additions and 0 deletions.
29 changes: 29 additions & 0 deletions tests/providers/google/cloud/hooks/test_bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -2206,6 +2206,35 @@ async def test_cancel_job_success(self, mock_job, mock_auth_default):

mock_job_instance.cancel.assert_called_once()

@pytest.mark.asyncio
@pytest.mark.db_test
@mock.patch("google.auth.default")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.Job")
async def test_cancel_job_failure(self, mock_job, mock_auth_default):
"""
Test that BigQueryAsyncHook handles exceptions during job cancellation correctly.
"""
mock_credentials = mock.MagicMock(spec=google.auth.compute_engine.Credentials)
mock_credentials.token = "ACCESS_TOKEN"
mock_auth_default.return_value = (mock_credentials, PROJECT_ID)

mock_job_instance = AsyncMock()
mock_job_instance.cancel.side_effect = Exception("Cancellation failed")
mock_job.return_value = mock_job_instance

hook = BigQueryAsyncHook()

job_id = "test_job_id"
project_id = "test_project"
location = "US"

with pytest.raises(Exception) as excinfo:
await hook.cancel_job(job_id=job_id, project_id=project_id, location=location)

assert "Cancellation failed" in str(excinfo.value), "Exception message not passed correctly"

mock_job_instance.cancel.assert_called_once()

@pytest.mark.asyncio
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.ClientSession")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_instance")
Expand Down

0 comments on commit 5371506

Please sign in to comment.