def test_set_yarn_spark_resource_config_fallback(
patched_virtual_memory,
patched_cpu_count,
patched_yarn_config,
patched_spark_config,
default_bootstrapper: Bootstrapper,
) -> None:
mocked_virtual_memory_total = PropertyMock(return_value=123 * 1024 * 1024)
type(patched_virtual_memory.return_value
).total = mocked_virtual_memory_total
patched_cpu_count.return_value = 456
default_bootstrapper.load_processing_job_config = MagicMock(
return_value=None)
default_bootstrapper.load_instance_type_info = MagicMock(return_value=None)
default_bootstrapper.get_yarn_spark_resource_config = MagicMock(
return_value=(patched_yarn_config, patched_spark_config))
default_bootstrapper.set_yarn_spark_resource_config()
patched_virtual_memory.assert_called_once()
mocked_virtual_memory_total.assert_called_once()
patched_cpu_count.assert_called_once()
default_bootstrapper.load_processing_job_config.assert_called_once()
default_bootstrapper.load_instance_type_info.assert_called_once()
default_bootstrapper.get_yarn_spark_resource_config.assert_called_once_with(
1, 123, 456)
patched_yarn_config.write_config.assert_called_once()
patched_spark_config.write_config.assert_called_once()
def test_load_instance_type_info(patched_exists, default_bootstrapper: Bootstrapper) -> None:
raw_config = [
{"InstanceType": "foo.xlarge", "foo": "bar"},
{"InstanceType": "bar.xlarge", "bar": "foo",},
]
exp_config = {"foo.xlarge": {"foo": "bar"}, "bar.xlarge": {"bar": "foo"}}
patched_exists.return_value = True
with patch("smspark.bootstrapper.open", mock_open(read_data=json.dumps(raw_config))) as m:
actual_config = default_bootstrapper.load_instance_type_info()
assert actual_config == exp_config
patched_exists.assert_called_once_with(Bootstrapper.INSTANCE_TYPE_INFO_PATH)
m.assert_called_once_with(Bootstrapper.INSTANCE_TYPE_INFO_PATH, "r")
def test_set_yarn_spark_resource_config(
patched_yarn_config, patched_spark_config, default_bootstrapper: Bootstrapper
) -> None:
processing_job_config = {
"ProcessingResources": {"ClusterConfig": {"InstanceType": "foo.xbar", "InstanceCount": 123}}
}
instance_type_info = {"foo.xbar": {"MemoryInfo": {"SizeInMiB": 456}, "VCpuInfo": {"DefaultVCpus": 789}}}
default_bootstrapper.load_processing_job_config = MagicMock(return_value=processing_job_config)
default_bootstrapper.load_instance_type_info = MagicMock(return_value=instance_type_info)
default_bootstrapper.get_yarn_spark_resource_config = MagicMock(
return_value=(patched_yarn_config, patched_spark_config)
)
default_bootstrapper.set_yarn_spark_resource_config()
default_bootstrapper.load_processing_job_config.assert_called_once()
default_bootstrapper.load_instance_type_info.assert_called_once()
default_bootstrapper.get_yarn_spark_resource_config.assert_called_once_with(123, 456, 789)
patched_yarn_config.write_config.assert_called_once()
patched_spark_config.write_config.assert_called_once()
def test_load_instance_type_info(patched_exists,
default_bootstrapper: Bootstrapper) -> None:
patched_exists.return_value = False
assert default_bootstrapper.load_instance_type_info() == {}
patched_exists.assert_called_once_with(
Bootstrapper.INSTANCE_TYPE_INFO_PATH)
