Repository: spotify/luigi
File: test/contrib/bigquery_test.py (Python)

Code before:

        job = MyBigQueryExtractTask()
        job.run()

        expected_body = {
            'configuration': {
                'extract': {
                    'sourceTable': {'projectId': 'proj', 'datasetId': 'ds', 'tableId': 't'},
                    'destinationUris': ['gs://_'],
                    'destinationFormat': 'AVRO',
                    'compression': 'NONE',
                    'useAvroLogicalTypes': True
                }
            }
        }
        run_job.assert_called_with('proj', expected_body, dataset=BQDataset('proj', 'ds', None))


class BigQueryClientTest(unittest.TestCase):

    def test_retry_succeeds_on_second_attempt(self):
        try:
            from googleapiclient import errors
        except ImportError:
            raise unittest.SkipTest('Unable to load googleapiclient module')
        client = MagicMock(spec=bigquery.BigQueryClient)
        attempts = 0

        @bigquery.bq_retry
        def fail_once(bq_client):
            nonlocal attempts
            attempts += 1
            if attempts == 1:
                raise errors.HttpError(
                    resp=MagicMock(status=500),
                    content=b'{"error": {"message": "stub"}',
                )
            else:
                return MagicMock(status=200)

        response = fail_once(client)
        client._initialise_client.assert_called_once()
        self.assertEqual(attempts, 2)
        self.assertEqual(response.status, 200)


Review comment (refactor):
Is this qualified name required because of mocking? Can we do
```python
from luigi.contrib.bigquery import BigQueryClient, BigQueryLoadTask, BigQueryTarget, BQDataset, \
    BigQueryRunQueryTask, BigQueryExtractTask
```

Code after:

        job = MyBigQueryExtractTask()
        job.run()

        expected_body = {
            'configuration': {
                'extract': {
                    'sourceTable': {'projectId': 'proj', 'datasetId': 'ds', 'tableId': 't'},
                    'destinationUris': ['gs://_'],
                    'destinationFormat': 'AVRO',
                    'compression': 'NONE',
                    'useAvroLogicalTypes': True
                }
            }
        }
        run_job.assert_called_with('proj', expected_body, dataset=BQDataset('proj', 'ds', None))


class BigQueryClientTest(unittest.TestCase):

    def test_retry_succeeds_on_second_attempt(self):
        try:
            from googleapiclient import errors
        except ImportError:
            raise unittest.SkipTest('Unable to load googleapiclient module')
        client = MagicMock(spec=BigQueryClient)
        attempts = 0

        @bigquery.bq_retry
        def fail_once(bq_client):
            nonlocal attempts
            attempts += 1
            if attempts == 1:
                raise errors.HttpError(
                    resp=MagicMock(status=500),
                    content=b'{"error": {"message": "stub"}',
                )
            else:
                return MagicMock(status=200)

        response = fail_once(client)
        client._initialise_client.assert_called_once()
        self.assertEqual(attempts, 2)
        self.assertEqual(response.status, 200)