diff --git a/cli/casp/src/casp/commands/init.py b/cli/casp/src/casp/commands/init.py index 7369816a9b..bd7e6fd2df 100644 --- a/cli/casp/src/casp/commands/init.py +++ b/cli/casp/src/casp/commands/init.py @@ -82,19 +82,30 @@ def _setup_custom_config(cfg: Dict[str, Any]): click.secho(f'Custom config path set to: {custom_config_path}', fg='green') -def _pull_image(): - """Pulls the docker image.""" - click.echo(f'Pulling Docker image: {docker_utils.DOCKER_IMAGE}...') - if not docker_utils.pull_image(): +def _pull_image_for_project(project: str = 'internal'): + """Pulls the docker image for the given project.""" + if not docker_utils.pull_image(docker_utils.PROJECT_TO_IMAGE[project]): click.secho( - f'\nError: Failed to pull Docker image {docker_utils.DOCKER_IMAGE}.', + (f'\nError: Failed to pull Docker image: ' + f'{docker_utils.PROJECT_TO_IMAGE[project]}.'), fg='red') click.secho('Initialization failed.', fg='red') sys.exit(1) @click.command(name='init', help='Initializes the CLI') -def cli(): +@click.option( + '--projects', + '--project', + '-p', + help=('The ClusterFuzz project to use. You can specify multiple projects.' + 'Ex.: -p dev -p internal'), + required=False, + default=('internal',), + type=click.Choice( + docker_utils.PROJECT_TO_IMAGE.keys(), case_sensitive=False), + multiple=True) +def cli(projects: tuple[str, ...]): """Initializes the CASP CLI. This is done by: @@ -117,5 +128,7 @@ def cli(): config.save_config(cfg) click.secho(f'Configuration saved to {config.CONFIG_FILE}.', fg='green') - _pull_image() + for project in projects: + _pull_image_for_project(project) + click.secho('Initialization complete.', fg='green') diff --git a/cli/casp/src/casp/tests/commands/test_init.py b/cli/casp/src/casp/tests/commands/test_init.py index be2a6d3437..8d20c0220f 100644 --- a/cli/casp/src/casp/tests/commands/test_init.py +++ b/cli/casp/src/casp/tests/commands/test_init.py @@ -45,7 +45,6 @@ def setUp(self): # Default mock behaviors for success paths self.mock_docker_utils.check_docker_setup.return_value = True self.mock_docker_utils.pull_image.return_value = True - self.mock_docker_utils.DOCKER_IMAGE = 'gcr.io/casp/runner:latest' credentials_path = '/fake/path/credentials.json' self.mock_gcloud.get_credentials_path.return_value = credentials_path self.mock_config.load_config.return_value = {} @@ -60,9 +59,6 @@ def test_init_success_all_steps(self): self.assertIn('Docker setup is correct.', result.output) self.assertIn('gcloud authentication is configured correctly.', result.output) - self.assertIn( - f'Pulling Docker image: {self.mock_docker_utils.DOCKER_IMAGE}', - result.output) self.assertIn('Initialization complete.', result.output) self.mock_docker_utils.check_docker_setup.assert_called_once() diff --git a/cli/casp/src/casp/tests/utils/test_docker.py b/cli/casp/src/casp/tests/utils/test_docker.py index 636dc3e52f..98fb599029 100644 --- a/cli/casp/src/casp/tests/utils/test_docker.py +++ b/cli/casp/src/casp/tests/utils/test_docker.py @@ -137,7 +137,7 @@ def test_pull_image_success(self, mock_echo, mock_secho, self.assertIn('Pulling Docker image:', args[0]) mock_check_docker_setup.assert_called_once() mock_images_collection.pull.assert_called_once_with( - docker_utils.DOCKER_IMAGE) + docker_utils.PROJECT_TO_IMAGE["internal"]) mock_secho.assert_not_called() @patch( @@ -176,12 +176,13 @@ def test_pull_image_not_found(self, mock_echo, mock_secho, self.assertFalse(result) mock_echo.assert_called_once_with( - f'Pulling Docker image: {docker_utils.DOCKER_IMAGE}...') + f'Pulling Docker image: {docker_utils.PROJECT_TO_IMAGE["internal"]}...') mock_check_docker_setup.assert_called_once() mock_images_collection.pull.assert_called_once_with( - docker_utils.DOCKER_IMAGE) - mock_secho.assert_called_once_with( - f'Error: Docker image {docker_utils.DOCKER_IMAGE} not found.', fg='red') + docker_utils.PROJECT_TO_IMAGE["internal"]) + mock_secho.assert_called_once() + args, _ = mock_secho.call_args + self.assertIn('not found', args[0]) if __name__ == '__main__': diff --git a/cli/casp/src/casp/utils/docker_utils.py b/cli/casp/src/casp/utils/docker_utils.py index d08f6c0e15..07cd6c8f0a 100644 --- a/cli/casp/src/casp/utils/docker_utils.py +++ b/cli/casp/src/casp/utils/docker_utils.py @@ -20,8 +20,15 @@ import docker # TODO: Make this configurable. -DOCKER_IMAGE = ("gcr.io/clusterfuzz-images/chromium/base/immutable/dev:" - "20251008165901-utc-893e97e-640142509185-compute-d609115-prod") +PROJECT_TO_IMAGE = { + 'dev': ("gcr.io/clusterfuzz-images/chromium/base/immutable/dev:" + "20251008165901-utc-893e97e-640142509185-compute-d609115-prod"), + 'internal': ( + "gcr.io/clusterfuzz-images/chromium/base/immutable/internal:" + "20251110132749-utc-363160d-640142509185-compute-c7f2f8c-prod"), + 'external': ("gcr.io/clusterfuzz-images/base/immutable/external:" + "20251111191918-utc-b5863ff-640142509185-compute-c5c296c-prod") +} def check_docker_setup() -> docker.client.DockerClient | None: @@ -52,16 +59,16 @@ def check_docker_setup() -> docker.client.DockerClient | None: return None -def pull_image() -> bool: +def pull_image(image: str = PROJECT_TO_IMAGE['internal']) -> bool: """Pulls the docker image.""" client = check_docker_setup() if not client: return False try: - click.echo(f'Pulling Docker image: {DOCKER_IMAGE}...') - client.images.pull(DOCKER_IMAGE) + click.echo(f'Pulling Docker image: {image}...') + client.images.pull(image) return True except docker.errors.DockerException: - click.secho(f'Error: Docker image {DOCKER_IMAGE} not found.', fg='red') + click.secho(f'Error: Docker image {image} not found.', fg='red') return False