# coding=utf-8 from unittest.mock import MagicMock, patch import pytest from common import config from spark_executor.core.spark_submit import ( SparkSubmitError, build_spark_submit_command, run_spark_submit, ) @pytest.fixture(autouse=True) def _restore_settings(): snapshot = config.Settings( data_dir=config.settings.data_dir, jobs_dir=config.settings.jobs_dir, log_dir=config.settings.log_dir, yarn_resource_manager_url=config.settings.yarn_resource_manager_url, log_level=config.settings.log_level, spark_submit_bin=config.settings.spark_submit_bin, ) yield config.settings.data_dir = snapshot.data_dir config.settings.jobs_dir = snapshot.jobs_dir config.settings.log_dir = snapshot.log_dir config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url config.settings.log_level = snapshot.log_level config.settings.spark_submit_bin = snapshot.spark_submit_bin def test_build_command_uses_provided_master_and_deploy_mode(): cmd = build_spark_submit_command( master="yarn", deploy_mode="cluster", script_path="/tmp/jobs/job_001.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, ) assert cmd[:2] == ["spark-submit", "--master"] # default binary assert "yarn" in cmd assert "cluster" in cmd assert "--queue" in cmd and "default" in cmd assert "--executor-memory" in cmd and "4G" in cmd assert "--executor-cores" in cmd and "2" in cmd assert "--num-executors" in cmd and "2" in cmd assert cmd[-1] == "/tmp/jobs/job_001.py" def test_extra_args_none_default(): cmd = build_spark_submit_command( master="yarn", deploy_mode="cluster", script_path="/tmp/jobs/job_001.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, ) # No extra flags between --conf block and script_path. assert cmd[-1] == "/tmp/jobs/job_001.py" def test_extra_args_single_pair(): cmd = build_spark_submit_command( master="yarn", deploy_mode="cluster", script_path="/tmp/jobs/job_001.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, extra_args={"jars": "hdfs:///libs/foo.jar"}, ) script_idx = cmd.index("/tmp/jobs/job_001.py") assert cmd[script_idx - 2 : script_idx] == [ "--jars", "hdfs:///libs/foo.jar", ] def test_extra_args_multiple_pairs_preserve_order(): cmd = build_spark_submit_command( master="yarn", deploy_mode="cluster", script_path="/tmp/jobs/job_001.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, extra_args={ "driver-memory": "2G", "name": "my-job", "py-files": "hdfs:///libs/udf.py", }, ) script_idx = cmd.index("/tmp/jobs/job_001.py") extras = cmd[script_idx - 6 : script_idx] assert extras == [ "--driver-memory", "2G", "--name", "my-job", "--py-files", "hdfs:///libs/udf.py", ] def test_extra_args_does_not_collide_with_spark_conf(): cmd = build_spark_submit_command( master="yarn", deploy_mode="cluster", script_path="/tmp/jobs/job_001.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, spark_conf={"spark.executor.memory": "8G"}, extra_args={"driver-memory": "2G"}, ) script_idx = cmd.index("/tmp/jobs/job_001.py") # --conf entries come before --extra_args entries. assert cmd[script_idx - 4 : script_idx] == [ "--conf", "spark.executor.memory=8G", "--driver-memory", "2G", ] def test_build_command_supports_standalone_master(): cmd = build_spark_submit_command( master="spark://master:7077", deploy_mode="client", script_path="/tmp/j.py", queue="default", executor_memory="1G", executor_cores=1, num_executors=1, ) assert "spark://master:7077" in cmd assert "client" in cmd def test_build_command_appends_spark_conf_entries(): cmd = build_spark_submit_command( master="yarn", deploy_mode="cluster", script_path="/tmp/j.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, spark_conf={"spark.sql.shuffle.partitions": "200", "spark.executor.memoryOverhead": "1G"}, ) # every k=v is emitted as --conf k=v assert "--conf" in cmd assert "spark.sql.shuffle.partitions=200" in cmd assert "spark.executor.memoryOverhead=1G" in cmd # script_path is still last assert cmd[-1] == "/tmp/j.py" def test_build_command_uses_settings_spark_submit_bin(): """The Spark CLI binary is env-configurable via settings.spark_submit_bin (env var SPARK_EXECUTOR_SPARK_SUBMIT_BIN). Default is 'spark-submit'; override to 'spark2-submit', 'pyspark', etc. for mixed-version hosts.""" config.settings.spark_submit_bin = "spark2-submit" cmd = build_spark_submit_command( master="yarn", deploy_mode="cluster", script_path="/tmp/j.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, ) assert cmd[0] == "spark2-submit" def test_build_command_supports_pyspark_binary(): config.settings.spark_submit_bin = "pyspark" cmd = build_spark_submit_command( master="yarn", deploy_mode="cluster", script_path="/tmp/j.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, ) assert cmd[0] == "pyspark" def test_run_spark_submit_returns_completed_process(monkeypatch): fake = MagicMock() fake.returncode = 0 fake.stderr = "Submitted application application_1\n" with patch("spark_executor.core.spark_submit.subprocess.run", return_value=fake) as m: result = run_spark_submit(["spark-submit", "/tmp/x.py"]) assert result is fake m.assert_called_once() # capture_output=True and text=True must be set kwargs = m.call_args.kwargs assert kwargs["capture_output"] is True assert kwargs["text"] is True def test_run_spark_submit_raises_on_nonzero_return(): fake = MagicMock() fake.returncode = 1 fake.stderr = "boom" with patch("spark_executor.core.spark_submit.subprocess.run", return_value=fake): with pytest.raises(SparkSubmitError, match="spark-submit failed"): run_spark_submit(["spark-submit", "/tmp/x.py"])