예제 #1
0
파일: bigquery_test.py 프로젝트: x35029/grr
    def testRetryUpload(self):
        bq_client = bigquery.BigQueryClient()

        resp = mock.Mock()
        resp.status = 503
        error = mock.Mock()
        error.resp = resp
        job = mock.Mock()
        # Always raise errors.HttpError on job.execute()
        job.configure_mock(
            **{"execute.side_effect": errors.HttpError(resp, b"nocontent")})
        job_id = "hunts_HFFE1D044_Results_1446056474"

        with tempfile.NamedTemporaryFile() as fd:
            fd.write("{data}")
            with mock.patch.object(time, "sleep") as mock_sleep:
                with self.assertRaises(bigquery.BigQueryJobUploadError):
                    bq_client.RetryUpload(job, job_id, error)

        # Make sure retry sleeps are correct.
        max_calls = config.CONFIG["BigQuery.retry_max_attempts"]
        retry_interval = config.CONFIG["BigQuery.retry_interval"]
        multiplier = config.CONFIG["BigQuery.retry_multiplier"]

        self.assertEqual(job.execute.call_count, max_calls)
        mock_sleep.assert_has_calls([
            mock.call(retry_interval),
            mock.call(retry_interval * multiplier)
        ])
예제 #2
0
    def testRetryUpload(self):
        bq_client = bigquery.BigQueryClient()

        resp = mock.Mock()
        resp.status = 503
        error = mock.Mock()
        error.resp = resp
        job = mock.Mock()
        # Always raise errors.HttpError on job.execute()
        job.configure_mock(
            **{"execute.side_effect": errors.HttpError(resp, b"nocontent")})
        job_id = "hunts_HFFE1D044_Results_1446056474"

        with temp.AutoTempFilePath() as filepath:
            with io.open(filepath, "w", encoding="utf-8") as filedesc:
                filedesc.write("{data}")

            with mock.patch.object(time, "sleep") as mock_sleep:
                with self.assertRaises(bigquery.BigQueryJobUploadError):
                    bq_client.RetryUpload(job, job_id, error)

        # Make sure retry sleeps are correct.
        max_calls = config.CONFIG["BigQuery.retry_max_attempts"]
        retry_interval = config.CONFIG["BigQuery.retry_interval"]
        multiplier = config.CONFIG["BigQuery.retry_multiplier"]

        self.assertEqual(job.execute.call_count, max_calls)
        mock_sleep.assert_has_calls([
            mock.call(retry_interval.ToFractional(rdfvalue.SECONDS)),
            mock.call(
                retry_interval.ToFractional(rdfvalue.SECONDS) * multiplier)
        ])