<s> """<STR_LIT>""" <EOL> import json <EOL> import os <EOL> import tempfile <EOL> import time <EOL> from apiclient . errors import HttpError <EOL> import mock <EOL> from grr . lib import bigquery <EOL> from grr . lib import config_lib <EOL> from grr . lib import flags <EOL> from grr . lib import rdfvalue <EOL> from grr . lib import test_lib <EOL> class BigQueryClientTest ( test_lib . GRRBaseTest ) : <EOL> """<STR_LIT>""" <EOL> TEST_KEY = """<STR_LIT>""" <EOL> PROJECT_ID = "<STR_LIT>" <EOL> SERVICE_ACCOUNT = "<STR_LIT>" <EOL> @ mock . patch . object ( bigquery , "<STR_LIT>" ) <EOL> @ mock . patch . object ( bigquery , "<STR_LIT>" ) <EOL> @ mock . patch . object ( bigquery . httplib2 , "<STR_LIT>" ) <EOL> def testInsertData ( self , mock_http , mock_build , mock_creds ) : <EOL> bq_client = bigquery . GetBigQueryClient ( <EOL> service_account = self . SERVICE_ACCOUNT , private_key = self . TEST_KEY , <EOL> project_id = self . PROJECT_ID ) <EOL> schema_data = json . load ( open ( os . path . join ( config_lib . CONFIG [ <EOL> "<STR_LIT>" ] , "<STR_LIT>" , "<STR_LIT>" ) ) ) <EOL> data_fd = open ( os . path . join ( config_lib . CONFIG [ <EOL> "<STR_LIT>" ] , "<STR_LIT>" , "<STR_LIT>" ) ) <EOL> now = rdfvalue . RDFDatetime ( ) . Now ( ) . AsSecondsFromEpoch ( ) <EOL> job_id = "<STR_LIT>" % now <EOL> bq_client . InsertData ( "<STR_LIT>" , data_fd , schema_data , job_id ) <EOL> insert = mock_build . return_value . jobs . return_value . insert <EOL> self . assertEqual ( insert . call_count , <NUM_LIT:1> ) <EOL> self . assertEqual ( job_id , insert . call_args_list [ <NUM_LIT:0> ] [ <NUM_LIT:1> ] [ <EOL> "<STR_LIT:body>" ] [ "<STR_LIT>" ] [ "<STR_LIT>" ] ) <EOL> def testRetryUpload ( self ) : <EOL> bq_client = bigquery . BigQueryClient ( ) <EOL> resp = mock . Mock ( ) <EOL> resp . status = <NUM_LIT> <EOL> error = mock . Mock ( ) <EOL> error . resp = resp <EOL> job = mock . Mock ( ) <EOL> job . configure_mock ( ** { "<STR_LIT>" : HttpError ( resp , "<STR_LIT>" ) } ) <EOL> job_id = "<STR_LIT>" <EOL> with tempfile . NamedTemporaryFile ( ) as fd : <EOL> fd . write ( "<STR_LIT>" ) <EOL> with mock . patch . object ( time , "<STR_LIT>" ) as mock_sleep : <EOL> with self . assertRaises ( bigquery . BigQueryJobUploadError ) : <EOL> bq_client . RetryUpload ( <EOL> job , job_id , error ) <EOL> max_calls = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> retry_interval = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> multiplier = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> self . assertEqual ( job . execute . call_count , max_calls ) <EOL> mock_sleep . assert_has_calls ( [ mock . call ( retry_interval ) , <EOL> mock . call ( retry_interval * multiplier ) ] ) <EOL> def main ( argv ) : <EOL> test_lib . GrrTestProgram ( argv = argv ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> flags . StartMain ( main ) </s>
