Added test for spark-run with pod identity

Author: Qmando

paasta_tools/cli/cmds/spark_run.py

@@ -431,6 +431,7 @@ def add_subparser(subparsers):
             " However, it must already be set for a different instance of "
             "the service."
         ),
+        default=None,
     )
 
     jupyter_group = list_parser.add_argument_group(
@@ -1237,7 +1238,6 @@ def paasta_spark_run(args: argparse.Namespace) -> int:
     if iam_role or args.force_pod_identity:
         if args.force_pod_identity:
             if args.yelpsoa_config_root != DEFAULT_SOA_DIR:
-                # TODO
                 print(
                     "--force-pod-identity cannot be used with --yelpsoa-config-root",
                     file=sys.stderr,

tests/cli/test_cmds_spark_run.py

@@ -1031,6 +1031,7 @@ def test_paasta_spark_run_bash(
         job_id=None,
         use_web_identity=False,
         uses_bulkdata=True,
+        force_pod_identity=False,
     )
     mock_load_system_paasta_config_utils.return_value.get_kube_clusters.return_value = (
         {}
@@ -1048,6 +1049,7 @@ def test_paasta_spark_run_bash(
     mock_spark_conf_builder.return_value.get_spark_conf.return_value = {
         "spark.kubernetes.executor.podTemplateFile": "/test/pod-template.yaml",
     }
+    mock_get_instance_config.return_value.get_iam_role.return_value = None
     spark_run.paasta_spark_run(args)
     mock_validate_work_dir.assert_called_once_with("/tmp/local")
     assert args.cmd == "/bin/bash"
@@ -1088,6 +1090,7 @@ def test_paasta_spark_run_bash(
         force_spark_resource_configs=False,
         use_eks=True,
         k8s_server_address=None,
+        service_account_name=None,
     )
     mock_spark_conf = mock_spark_conf_builder.return_value.get_spark_conf.return_value
     mock_configure_and_run_docker_container.assert_called_once_with(
@@ -1158,6 +1161,7 @@ def test_paasta_spark_run(
         job_id=None,
         use_web_identity=False,
         uses_bulkdata=True,
+        force_pod_identity=False,
     )
     mock_load_system_paasta_config_utils.return_value.get_kube_clusters.return_value = (
         {}
@@ -1178,6 +1182,7 @@ def test_paasta_spark_run(
     mock_spark_conf_builder.return_value.get_spark_conf.return_value = {
         "spark.kubernetes.executor.podTemplateFile": "/test/pod-template.yaml",
     }
+    mock_get_instance_config.return_value.get_iam_role.return_value = None
     spark_run.paasta_spark_run(args)
     mock_validate_work_dir.assert_called_once_with("/tmp/local")
     assert args.cmd == "USER=test timeout 1m spark-submit test.py"
@@ -1218,6 +1223,7 @@ def test_paasta_spark_run(
         force_spark_resource_configs=False,
         use_eks=True,
         k8s_server_address=None,
+        service_account_name=None,
     )
     mock_configure_and_run_docker_container.assert_called_once_with(
         args,
@@ -1285,6 +1291,7 @@ def test_paasta_spark_run_pyspark(
         job_id=None,
         use_web_identity=False,
         uses_bulkdata=True,
+        force_pod_identity=False,
     )
     mock_load_system_paasta_config_utils.return_value.get_kube_clusters.return_value = (
         {}
@@ -1306,6 +1313,7 @@ def test_paasta_spark_run_pyspark(
     mock_spark_conf_builder.return_value.get_spark_conf.return_value = {
         "spark.kubernetes.executor.podTemplateFile": "/test/pod-template.yaml",
     }
+    mock_get_instance_config.return_value.get_iam_role.return_value = None
 
     spark_run.paasta_spark_run(args)
     mock_validate_work_dir.assert_called_once_with("/tmp/local")
@@ -1347,6 +1355,7 @@ def test_paasta_spark_run_pyspark(
         force_spark_resource_configs=False,
         use_eks=True,
         k8s_server_address=None,
+        service_account_name=None,
     )
     mock_configure_and_run_docker_container.assert_called_once_with(
         args,
@@ -1428,6 +1437,7 @@ def test_paasta_spark_run_uses_bulkdata(
         job_id=None,
         use_web_identity=False,
         uses_bulkdata=spark_run_arg_uses_bulkdata,
+        force_pod_identity=False,
     )
     mock_load_system_paasta_config_spark_run.return_value.get_pools_for_cluster.return_value = [
         "test-pool"
@@ -1436,6 +1446,7 @@ def test_paasta_spark_run_uses_bulkdata(
     mock_get_instance_config.return_value.config_dict = {
         "uses_bulkdata": instance_config_uses_bulkdata
     }
+    mock_get_instance_config.return_value.get_iam_role.return_value = None
 
     spark_run.paasta_spark_run(args)
 
@@ -1644,3 +1655,179 @@ def test_get_aws_credentials_default_profile(mock_boto3_session, mock_use_aws_pr
     )
 
     mock_use_aws_profile.assert_called_once_with()
+
+
+@mock.patch.object(spark_run, "get_all_iam_roles_for_service", autospec=True)
+@mock.patch.object(spark_run, "validate_work_dir", autospec=True)
+@mock.patch.object(utils, "load_system_paasta_config", autospec=True)
+@mock.patch.object(spark_run, "load_system_paasta_config", autospec=True)
+@mock.patch.object(spark_run, "get_instance_config", autospec=True)
+@mock.patch.object(spark_run, "get_aws_credentials", autospec=True)
+@mock.patch.object(spark_run, "get_docker_image", autospec=True)
+@mock.patch.object(spark_run, "get_spark_app_name", autospec=True)
+@mock.patch.object(spark_run, "auto_add_timeout_for_spark_job", autospec=True)
+@mock.patch.object(spark_run, "_parse_user_spark_args", autospec=True)
+@mock.patch(
+    "paasta_tools.cli.cmds.spark_run.spark_config.SparkConfBuilder", autospec=True
+)
+@mock.patch.object(spark_run, "configure_and_run_docker_container", autospec=True)
+@mock.patch.object(spark_run, "get_smart_paasta_instance_name", autospec=True)
+@pytest.mark.parametrize(
+    "iam_role_set,force_pod_identity,aws_creds_provided,custom_yelpsoa,should_exit_early",
+    [
+        # Service has iam_role set, no overrides
+        ("arn:aws:iam::123456789:role/test", False, None, None, 0),
+        # Service has iam_role set, you can override aws_creds_yamls
+        ("arn:aws:iam::123456789:role/test", False, True, None, 0),
+        # Force-pod-identity set, should fail as it's not in the (empty) allowed set
+        (None, "arn:aws:iam::123456789:role/test", None, None, 1),
+        # Force-pod-identity set, succeeds because test adds it to allowed set
+        (None, "arn:aws:iam::123456789:role/test", None, None, 0),
+        # Force-pod-identity set, different from set iam_role, fails as it's not in allowed set
+        (
+            "arn:aws:iam::123456789:role/test",
+            "arn:aws:iam::123456789:role/test2",
+            None,
+            None,
+            1,
+        ),
+        # Force-pod-identity set, but fails because custom soa dir provided
+        (None, "arn:aws:iam::123456789:role/test", None, "/custom/soa", 1),
+    ],
+)
+def test_paasta_spark_run_with_pod_identity(
+    mock_get_smart_paasta_instance_name,
+    mock_configure_and_run_docker_container,
+    mock_spark_conf_builder,
+    mock_parse_user_spark_args,
+    mock_auto_add_timeout_for_spark_job,
+    mock_get_spark_app_name,
+    mock_get_docker_image,
+    mock_get_aws_credentials,
+    mock_get_instance_config,
+    mock_load_system_paasta_config_spark_run,
+    mock_load_system_paasta_config_utils,
+    mock_validate_work_dir,
+    mock_get_all_iam_roles_for_service,
+    iam_role_set,
+    force_pod_identity,
+    aws_creds_provided,
+    custom_yelpsoa,
+    should_exit_early,
+):
+    args = argparse.Namespace(
+        work_dir="/tmp/local",
+        cmd="USER=test spark-submit test.py",
+        build=True,
+        image=None,
+        enable_compact_bin_packing=False,
+        disable_compact_bin_packing=False,
+        service="test-service",
+        instance="test-instance",
+        cluster="test-cluster",
+        pool="test-pool",
+        yelpsoa_config_root=custom_yelpsoa or spark_run.DEFAULT_SOA_DIR,
+        aws_credentials_yaml="/path/to/creds" if aws_creds_provided else None,
+        aws_profile=None,
+        spark_args="spark.cores.max=100 spark.executor.cores=10",
+        cluster_manager=spark_run.CLUSTER_MANAGER_K8S,
+        timeout_job_runtime="1m",
+        enable_dra=True,
+        aws_region="test-region",
+        force_spark_resource_configs=False,
+        assume_aws_role=None,
+        aws_role_duration=3600,
+        k8s_server_address=None,
+        tronfig=None,
+        job_id=None,
+        use_web_identity=False,
+        uses_bulkdata=True,
+        force_pod_identity=force_pod_identity,
+    )
+    # Use the expected return code to set the allowed iam roles
+    if force_pod_identity and not should_exit_early:
+        mock_get_all_iam_roles_for_service.return_value = (force_pod_identity,)
+    else:
+        mock_get_all_iam_roles_for_service.return_value = (
+            "arn:aws:iam::123456789:role/unreferenced-role",
+        )
+    mock_load_system_paasta_config_utils.return_value.get_kube_clusters.return_value = (
+        {}
+    )
+    mock_load_system_paasta_config_spark_run.return_value.get_cluster_aliases.return_value = (
+        {}
+    )
+    mock_load_system_paasta_config_spark_run.return_value.get_pools_for_cluster.return_value = [
+        "test-pool"
+    ]
+    mock_load_system_paasta_config_spark_run.return_value.get_eks_cluster_aliases.return_value = {
+        "test-cluster": "test-cluster"
+    }
+    mock_get_docker_image.return_value = DUMMY_DOCKER_IMAGE_DIGEST
+    mock_auto_add_timeout_for_spark_job.return_value = (
+        "USER=test timeout 1m spark-submit test.py"
+    )
+    mock_spark_conf_builder.return_value.get_spark_conf.return_value = {
+        "spark.kubernetes.executor.podTemplateFile": "/test/pod-template.yaml",
+    }
+    mock_get_instance_config.return_value.get_iam_role.return_value = (
+        "arn:aws:iam::123456789:role/test"
+    )
+    ret = spark_run.paasta_spark_run(args)
+    if should_exit_early:
+        assert ret != 0
+        return
+    mock_validate_work_dir.assert_called_once_with("/tmp/local")
+    assert args.cmd == "USER=test timeout 1m spark-submit test.py"
+    mock_get_instance_config.assert_called_once_with(
+        service="test-service",
+        instance="test-instance",
+        cluster="test-cluster",
+        load_deployments=False,
+        soa_dir=spark_run.DEFAULT_SOA_DIR,
+    )
+    mock_get_aws_credentials.assert_called_once_with(
+        service="test-service",
+        aws_credentials_yaml="/path/to/creds"
+        if aws_creds_provided
+        else "/etc/boto_cfg/mrjob.yaml",
+        profile_name=None,
+        assume_aws_role_arn=None,
+        session_duration=3600,
+        use_web_identity=False,
+    )
+    mock_get_docker_image.assert_called_once_with(
+        args, mock_get_instance_config.return_value
+    )
+    mock_get_spark_app_name.assert_called_once_with("USER=test spark-submit test.py")
+    mock_parse_user_spark_args.assert_called_once_with(
+        "spark.cores.max=100 spark.executor.cores=10"
+    )
+    mock_spark_conf_builder.return_value.get_spark_conf.assert_called_once_with(
+        cluster_manager=spark_run.CLUSTER_MANAGER_K8S,
+        spark_app_base_name=mock_get_spark_app_name.return_value,
+        docker_img=DUMMY_DOCKER_IMAGE_DIGEST,
+        user_spark_opts=mock_parse_user_spark_args.return_value,
+        paasta_cluster="test-cluster",
+        paasta_pool="test-pool",
+        paasta_service="test-service",
+        paasta_instance=mock_get_smart_paasta_instance_name.return_value,
+        extra_volumes=mock_get_instance_config.return_value.get_volumes.return_value,
+        aws_creds=None,
+        aws_region="test-region",
+        force_spark_resource_configs=False,
+        use_eks=True,
+        k8s_server_address=None,
+        service_account_name="paasta--arn-aws-iam-123456789-role-test",
+    )
+    mock_configure_and_run_docker_container.assert_called_once_with(
+        args,
+        docker_img=DUMMY_DOCKER_IMAGE_DIGEST,
+        instance_config=mock_get_instance_config.return_value,
+        system_paasta_config=mock_load_system_paasta_config_spark_run.return_value,
+        spark_conf=mock_spark_conf_builder.return_value.get_spark_conf.return_value,
+        aws_creds=mock_get_aws_credentials.return_value,
+        cluster_manager=spark_run.CLUSTER_MANAGER_K8S,
+        pod_template_path="/test/pod-template.yaml",
+        extra_driver_envs=dict(),
+    )