def test_initialize_queues(self):
# Act
with metadata.local_project_id_override('my-project-id'):
google_cloud_task_queue_config.initialize_queues(
google_auth_token='fake-auth-token')
# Assert
queues_updated_by_id = self.get_updated_queues()
for queue in queues_updated_by_id.values():
self.assertTrue(
queue.name.startswith(
'projects/my-project-id/locations/us-east1/queues/'))
self.assertEqual(queue.stackdriver_logging_config.sampling_ratio,
1.0)
direct_ingest_queue_ids = {
'direct-ingest-state-process-job-queue-v2',
'direct-ingest-jpp-process-job-queue-v2',
'direct-ingest-bq-import-export-v2', 'direct-ingest-scheduler-v2'
}
self.assertFalse(
direct_ingest_queue_ids.difference(queues_updated_by_id.keys()))
for queue_id in direct_ingest_queue_ids:
queue = queues_updated_by_id[queue_id]
self.assertEqual(queue.rate_limits.max_concurrent_dispatches, 1)
for region in regions.get_supported_regions():
self.assertTrue(region.get_queue_name() in queues_updated_by_id)
self.assertTrue('bigquery-v2' in queues_updated_by_id)
self.assertTrue('job-monitor-v2' in queues_updated_by_id)
self.assertTrue('scraper-phase-v2' in queues_updated_by_id)
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--project_id',
required=True,
help='Project to initialize queues for')
parser.add_argument('--google_auth_token',
required=True,
help='Auth token (obtained via '
'`gcloud auth print-access-token`).')
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format='%(message)s')
google_cloud_task_queue_config.initialize_queues(
google_auth_token=args.google_auth_token, project_id=args.project_id)
def test_initialize_queues(self):
# Act
google_cloud_task_queue_config.initialize_queues(
google_auth_token='fake-auth-token', project_id='my-project-id')
# Assert
queues_updated_by_id: Dict[str, queue_pb2.Queue] = {}
for method_name, args, _kwargs in self.mock_client.mock_calls:
if method_name == 'update_queue':
queue = args[0]
if not isinstance(queue, queue_pb2.Queue):
self.fail(f"Unexpected type [{type(queue)}]")
_, queue_id = os.path.split(queue.name)
queues_updated_by_id[queue_id] = queue
for queue in queues_updated_by_id.values():
self.assertTrue(
queue.name.startswith(
'projects/my-project-id/locations/us-east1/queues/'))
self.assertEqual(queue.stackdriver_logging_config.sampling_ratio,
1.0)
direct_ingest_queue_ids = {
'direct-ingest-state-process-job-queue-v2',
'direct-ingest-jpp-process-job-queue-v2',
'direct-ingest-scheduler-v2'
}
self.assertFalse(
direct_ingest_queue_ids.difference(queues_updated_by_id.keys()))
for queue_id in direct_ingest_queue_ids:
queue = queues_updated_by_id[queue_id]
self.assertEqual(queue.rate_limits.max_concurrent_dispatches, 1)
for region in regions.get_supported_regions():
self.assertTrue(region.get_queue_name() in queues_updated_by_id)
self.assertTrue('bigquery-v2' in queues_updated_by_id)
self.assertTrue('job-monitor-v2' in queues_updated_by_id)
self.assertTrue('scraper-phase-v2' in queues_updated_by_id)
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
# =============================================================================
"""Script run on deploy that initializes all task queues with appropriate
configurations."""
import argparse
import logging
from recidiviz.utils import metadata
from recidiviz.common.google_cloud import google_cloud_task_queue_config
if __name__ == '__main__':
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--project_id',
required=True,
help='Project to initialize queues for')
parser.add_argument('--google_auth_token',
required=True,
help='Auth token (obtained via '
'`gcloud auth print-access-token`).')
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format='%(message)s')
with metadata.local_project_id_override(args.project_id):
google_cloud_task_queue_config.initialize_queues(
google_auth_token=args.google_auth_token, )
def test_initialize_queues(self, mock_regions):
# Arrange
region_xx = fake_region(
region_code='us_xx',
queue={'rate_limits': {
'max_dispatches_per_second': 0.3
}})
region_xx.get_queue_name.return_value = 'us_xx_queue'
region_yy = fake_region(region_code='us_yy')
region_yy.get_queue_name.return_value = 'us_yy_queue'
mock_regions.return_value = [region_xx, region_yy]
# Act
with metadata.local_project_id_override('my-project-id'):
google_cloud_task_queue_config.initialize_queues(
google_auth_token='fake-auth-token')
# Assert
queues_updated_by_id = self.get_updated_queues()
for queue in queues_updated_by_id.values():
self.assertTrue(
queue.name.startswith(
'projects/my-project-id/locations/us-east1/queues/'))
self.assertEqual(queue.stackdriver_logging_config.sampling_ratio,
1.0)
direct_ingest_queue_ids = {
'direct-ingest-state-process-job-queue-v2',
'direct-ingest-jpp-process-job-queue-v2',
'direct-ingest-bq-import-export-v2', 'direct-ingest-scheduler-v2'
}
self.assertFalse(
direct_ingest_queue_ids.difference(queues_updated_by_id.keys()))
for queue_id in direct_ingest_queue_ids:
queue = queues_updated_by_id[queue_id]
self.assertEqual(queue.rate_limits.max_concurrent_dispatches, 1)
# Test that composition works as expected
self.assertEqual(
queues_updated_by_id[region_xx.get_queue_name()],
queue_pb2.Queue(
name=
'projects/my-project-id/locations/us-east1/queues/us_xx_queue',
rate_limits=queue_pb2.RateLimits(
# This is overridden in the mock above
max_dispatches_per_second=0.3,
max_concurrent_dispatches=3,
),
retry_config=queue_pb2.RetryConfig(
min_backoff=duration_pb2.Duration(seconds=5),
max_backoff=duration_pb2.Duration(seconds=300),
max_attempts=5,
),
stackdriver_logging_config=queue_pb2.StackdriverLoggingConfig(
sampling_ratio=1.0, )))
# Test that other regions are unaffected
self.assertEqual(
queues_updated_by_id[region_yy.get_queue_name()],
queue_pb2.Queue(
name=
'projects/my-project-id/locations/us-east1/queues/us_yy_queue',
rate_limits=queue_pb2.RateLimits(
max_dispatches_per_second=0.08333333333,
max_concurrent_dispatches=3,
),
retry_config=queue_pb2.RetryConfig(
min_backoff=duration_pb2.Duration(seconds=5),
max_backoff=duration_pb2.Duration(seconds=300),
max_attempts=5,
),
stackdriver_logging_config=queue_pb2.StackdriverLoggingConfig(
sampling_ratio=1.0, )))
self.assertTrue('bigquery-v2' in queues_updated_by_id)
self.assertTrue('job-monitor-v2' in queues_updated_by_id)
self.assertTrue('scraper-phase-v2' in queues_updated_by_id)
def test_initialize_queues(self, mock_regions: Mock) -> None:
# Arrange
region_xx = fake_region(
region_code="us_xx",
queue={"rate_limits": {
"max_dispatches_per_second": 0.3
}},
get_queue_name=lambda: "us_xx_queue",
)
region_yy = fake_region(region_code="us_yy",
get_queue_name=lambda: "us_yy_queue")
mock_regions.return_value = [region_xx, region_yy]
# Act
with metadata.local_project_id_override("my-project-id"):
google_cloud_task_queue_config.initialize_queues(
google_auth_token="fake-auth-token")
# Assert
queues_updated_by_id = self.get_updated_queues()
for queue in queues_updated_by_id.values():
self.assertTrue(
queue.name.startswith(
"projects/my-project-id/locations/us-east1/queues/"))
self.assertEqual(queue.stackdriver_logging_config.sampling_ratio,
1.0)
direct_ingest_queue_ids = {
"direct-ingest-state-process-job-queue-v2",
"direct-ingest-jpp-process-job-queue-v2",
"direct-ingest-bq-import-export-v2",
"direct-ingest-scheduler-v2",
}
self.assertFalse(
direct_ingest_queue_ids.difference(queues_updated_by_id.keys()))
for queue_id in direct_ingest_queue_ids:
queue = queues_updated_by_id[queue_id]
self.assertEqual(queue.rate_limits.max_concurrent_dispatches, 1)
# Test that composition works as expected
self.assertEqual(
queues_updated_by_id[region_xx.get_queue_name()],
queue_pb2.Queue(
name=
"projects/my-project-id/locations/us-east1/queues/us_xx_queue",
rate_limits=queue_pb2.RateLimits(
# This is overridden in the mock above
max_dispatches_per_second=0.3,
max_concurrent_dispatches=3,
),
retry_config=queue_pb2.RetryConfig(
min_backoff=duration_pb2.Duration(seconds=5),
max_backoff=duration_pb2.Duration(seconds=300),
max_attempts=5,
),
stackdriver_logging_config=queue_pb2.StackdriverLoggingConfig(
sampling_ratio=1.0, ),
),
)
# Test that other regions are unaffected
self.assertEqual(
queues_updated_by_id[region_yy.get_queue_name()],
queue_pb2.Queue(
name=
"projects/my-project-id/locations/us-east1/queues/us_yy_queue",
rate_limits=queue_pb2.RateLimits(
max_dispatches_per_second=0.08333333333,
max_concurrent_dispatches=3,
),
retry_config=queue_pb2.RetryConfig(
min_backoff=duration_pb2.Duration(seconds=5),
max_backoff=duration_pb2.Duration(seconds=300),
max_attempts=5,
),
stackdriver_logging_config=queue_pb2.StackdriverLoggingConfig(
sampling_ratio=1.0, ),
),
)
self.assertTrue("bigquery-v2" in queues_updated_by_id)
self.assertTrue("job-monitor-v2" in queues_updated_by_id)
self.assertTrue("scraper-phase-v2" in queues_updated_by_id)
self.assertTrue("admin-panel-data-discovery" in queues_updated_by_id)