[CLI] feat: add image selection and fix tests

cli/casp/src/casp/commands/init.py

@@ -17,6 +17,7 @@
 import sys
 from typing import Any
 from typing import Dict
+from typing import Tuple
 
 from casp.utils import config
 from casp.utils import docker_utils
@@ -82,19 +83,28 @@ def _setup_custom_config(cfg: Dict[str, Any]):
   click.secho(f'Custom config path set to: {custom_config_path}', fg='green')
 
 
-def _pull_image():
+def _pull_image(image: str = 'internal'):
   """Pulls the docker image."""
-  click.echo(f'Pulling Docker image: {docker_utils.DOCKER_IMAGE}...')
-  if not docker_utils.pull_image():
+  if not docker_utils.pull_image(image):
     click.secho(
-        f'\nError: Failed to pull Docker image {docker_utils.DOCKER_IMAGE}.',
+        (f'\nError: Failed to pull Docker image: '
+         f'{docker_utils.DOCKER_IMAGES[image]}.'),
         fg='red')
     click.secho('Initialization failed.', fg='red')
     sys.exit(1)
 
 
 @click.command(name='init', help='Initializes the CLI')
-def cli():
+@click.option(
+    '--image',
+    '-i',
+    help=('The Docker image to use. You can specify multiple images.'
+          'Ex.: --image dev --image internal'),
+    required=False,
+    default=('internal',),
+    type=click.Choice(['dev', 'internal', 'external'], case_sensitive=False),
+    multiple=True)
+def cli(image: Tuple[str, ...]):
   """Initializes the CASP CLI.
 
   This is done by:
@@ -117,5 +127,7 @@ def cli():
   config.save_config(cfg)
   click.secho(f'Configuration saved to {config.CONFIG_FILE}.', fg='green')
 
-  _pull_image()
+  for img in image:
+    _pull_image(img)
+
   click.secho('Initialization complete.', fg='green')

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()

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.DOCKER_IMAGES["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.DOCKER_IMAGES["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.DOCKER_IMAGES["internal"])
+    mock_secho.assert_called_once()
+    args, _ = mock_secho.call_args
+    self.assertIn('not found', args[0])
 
 
 if __name__ == '__main__':

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")
+DOCKER_IMAGES = {
+    '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,17 @@ def check_docker_setup() -> docker.client.DockerClient | None:
     return None
 
 
-def pull_image() -> bool:
+def pull_image(image: str = '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: {DOCKER_IMAGES[image]}...')
+    client.images.pull(DOCKER_IMAGES[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 {DOCKER_IMAGES[image]} not found.', fg='red')
     return False