Updates tests to include location.

google-marketing-solutions · Aug 16, 2024 · a904add · a904add
1 parent 4fac9cc
commit a904add
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 12 deletions.
diff --git a/backend/tests/jobs/unit/workers/bq_waiter_tests.py b/backend/tests/jobs/unit/workers/bq_waiter_tests.py
@@ -29,10 +29,11 @@ def test_execute_job_with_status(self, job_status, enqueue_called):
         'logger_credentials': _make_credentials(),
     }
     worker_inst = bq_waiter.BQWaiter(
-        {'job_id': 'JOBID',}, 1, 1, **logging_creds)
+        {'job_id': 'JOBID', 'location': 'US'}, 1, 1, **logging_creds)
     mock_job = mock.create_autospec(
         bigquery.job.QueryJob, instance=True, spec_set=True)
     mock_job.job_id = 'JOBID'
+    mock_job.location = 'US'
     mock_job.error_result = None
     mock_job.state = job_status
     mock_job.done.return_value = not enqueue_called
@@ -54,10 +55,11 @@ def test_execute_job_with_status(self, job_status, enqueue_called):
             spec_set=True))
     self.enter_context(mock.patch.object(worker_inst, '_log', autospec=True))
     worker_inst._execute()
+    mock_client.get_job.assert_called_with(job_id='JOBID', location='US')
     if enqueue_called:
       patched_enqueue.assert_called_once()
       self.assertEqual(patched_enqueue.call_args[0][0], 'BQWaiter')
-      self.assertEqual(patched_enqueue.call_args[0][1], {'job_id': 'JOBID'})
+      self.assertEqual(patched_enqueue.call_args[0][1], {'job_id': 'JOBID', 'location': 'US'})
     else:
       patched_enqueue.assert_not_called()
 
@@ -78,7 +80,7 @@ def test_job_error_raises_worker_exception(self):
               autospec=True,
               spec_set=True))
       worker_inst = bq_waiter.BQWaiter(
-          {'job_id': 'JOBID',}, 1, 1,
+          {'job_id': 'JOBID', 'location': 'US'}, 1, 1,
           logger_project='PROJECT',
           logger_credentials=_make_credentials())
       worker_inst._execute()

diff --git a/backend/tests/jobs/unit/workers/bq_worker_tests.py b/backend/tests/jobs/unit/workers/bq_worker_tests.py
@@ -40,24 +40,24 @@ def test_enqueues_bqwaiter_after_some_time(self, job_is_done, enqueue_called):
     }
     worker_inst = bq_worker.BQWorker(
         {'bq_project_id': 'BQID'}, 1, 1, **logging_creds)
-    mock_client = mock.create_autospec(
-        bigquery.Client, instance=True, spec_set=False)
-    mock_client.project = 'PROJECT'
-    job = bigquery.job.QueryJob('JOBID', 'query', mock_client)
-    self.enter_context(
-        mock.patch.object(job, 'done', return_value=job_is_done, autospec=True))
+    mock_job = mock.create_autospec(
+        bigquery.job.QueryJob, instance=True, spec_set=False)
+    mock_job.done.return_value = job_is_done
+    mock_job.job_id = 'JOBID'
+    mock_job.location = 'US'
+    mock_job.error_result = None
     patched_enqueue = self.enter_context(
         mock.patch.object(
             worker_inst,
             '_enqueue',
             return_value=True,
             autospec=True,
             spec_set=True))
-    worker_inst._wait(job)
+    worker_inst._wait(mock_job)
     if enqueue_called:
       patched_enqueue.assert_called_once()
       self.assertEqual(patched_enqueue.call_args[0][0], 'BQWaiter')
-      self.assertEqual(patched_enqueue.call_args[0][1], {'job_id': 'JOBID'})
+      self.assertEqual(patched_enqueue.call_args[0][1], {'job_id': 'JOBID', 'location': 'US'})
     else:
       patched_enqueue.assert_not_called()
 
@@ -93,7 +93,7 @@ class BQWorkerGetClientTest(parameterized.TestCase):
         'client_info_user_agent': 'cloud-solutions/crmint-usage-v3',
       },
       {
-        'report_usage_id_present': False, 
+        'report_usage_id_present': False,
         'client_info_user_agent': None
       },
   )