From 31673fc21e8b4a0e5b7177829100f84652d8c796 Mon Sep 17 00:00:00 2001 From: Luis Perez Date: Wed, 23 Oct 2024 11:06:15 -0700 Subject: [PATCH 1/5] Add tron topology_stpread_constraints support to PaaSTA This adds support for configuring Tron-launched pods with a default Topology Spread Constraint (and node affinities) that will spread pods out across multiple AZs - otherwise, Karpenter will overwhelmingly favor a single AZ due to our config --- paasta_tools/tron_tools.py | 43 ++++++++++++++++++++++++++++++++++++-- tests/test_tron_tools.py | 2 +- 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/paasta_tools/tron_tools.py b/paasta_tools/tron_tools.py index ccccf5f557..f5918fce25 100644 --- a/paasta_tools/tron_tools.py +++ b/paasta_tools/tron_tools.py @@ -62,7 +62,9 @@ from paasta_tools import spark_tools from paasta_tools.kubernetes_tools import ( + NodeSelectorConfig, allowlist_denylist_to_requirements, + contains_zone_label, get_service_account_name, limit_size_with_hash, raw_selectors_to_requirements, @@ -248,6 +250,7 @@ class TronActionConfigDict(InstanceConfigDict, total=False): # maneuvering to unify command: str service_account_name: str + node_selectors: Dict[str, NodeSelectorConfig] # the values for this dict can be anything since it's whatever # spark accepts @@ -594,18 +597,35 @@ def get_node_selectors(self) -> Dict[str, str]: def get_node_affinities(self) -> Optional[List[Dict[str, Union[str, List[str]]]]]: """Converts deploy_whitelist and deploy_blacklist in node affinities. - note: At the time of writing, `kubectl describe` does not show affinities, + NOTE: At the time of writing, `kubectl describe` does not show affinities, only selectors. To see affinities, use `kubectl get pod -o json` instead. + + WARNING: At the time of writing, we only used requiredDuringSchedulingIgnoredDuringExecution node affinities in Tron as we currently have + no use case for preferredDuringSchedulingIgnoredDuringExecution node affinities. """ requirements = allowlist_denylist_to_requirements( allowlist=self.get_deploy_whitelist(), denylist=self.get_deploy_blacklist(), ) + node_selectors = self.config_dict.get("node_selectors", {}) requirements.extend( raw_selectors_to_requirements( - raw_selectors=self.config_dict.get("node_selectors", {}), # type: ignore + raw_selectors=self.config_dict.get("node_selectors", {}), ) ) + + # PAASTA-18198: To improve AZ balance with Karpenter, we temporarily allow specifying zone affinities per pool + pool_node_affinities = load_system_paasta_config().get_pool_node_affinities() + if pool_node_affinities and self.get_pool() in pool_node_affinities: + current_pool_node_affinities = pool_node_affinities[self.get_pool()] + # If the service already has a node selector for a zone, we don't want to override it + if current_pool_node_affinities and not contains_zone_label(node_selectors): + requirements.extend( + raw_selectors_to_requirements( + raw_selectors=current_pool_node_affinities, + ) + ) + if not requirements: return None @@ -984,6 +1004,25 @@ def format_tron_action_dict(action_config: TronActionConfig): result["node_selectors"] = action_config.get_node_selectors() result["node_affinities"] = action_config.get_node_affinities() + # XXX: this is currently hardcoded since we should only really need TSC for zone-aware scheduling + result["topology_spread_constraints"] = [ + { + # try to evenly spread pods across specified topology + "max_skew": 1, + # narrow down what pods to consider when spreading + "label_selector": { + # only consider pods that are managed by tron + "app.kubernetes.io/managed-by": "tron", + # and in the same pool + "paasta.yelp.com/pool": action_config.get_pool(), + }, + # now, spread across AZs + "topology_key": "topology.kubernetes.io/zone", + # but if not possible, schedule even with a zonal imbalance + "when_unsatisfiable": "ScheduleAnyway", + }, + ] + # XXX: once we're off mesos we can make get_cap_* return just the cap names as a list result["cap_add"] = [cap["value"] for cap in action_config.get_cap_add()] result["cap_drop"] = [cap["value"] for cap in action_config.get_cap_drop()] diff --git a/tests/test_tron_tools.py b/tests/test_tron_tools.py index aea121dd06..6067a9a9a1 100644 --- a/tests/test_tron_tools.py +++ b/tests/test_tron_tools.py @@ -1896,7 +1896,7 @@ def test_create_complete_config_e2e(self, tmpdir): # that are not static, this will cause continuous reconfiguration, which # will add significant load to the Tron API, which happened in DAR-1461. # but if this is intended, just change the hash. - assert hasher.hexdigest() == "0585c2049596190f5510f8ce4fc3a9ac" + assert hasher.hexdigest() == "31634ab048abe9b40b71851797d48e4d" def test_override_default_pool_override(self, tmpdir): soa_dir = tmpdir.mkdir("test_create_complete_config_soa") From c1f62271866fe41bd38c15ca6d9f9f491d7a49fe Mon Sep 17 00:00:00 2001 From: Luis Perez Date: Thu, 24 Oct 2024 09:25:26 -0700 Subject: [PATCH 2/5] Update tests --- tests/test_tron_tools.py | 55 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/tests/test_tron_tools.py b/tests/test_tron_tools.py index 6067a9a9a1..7e9fb54000 100644 --- a/tests/test_tron_tools.py +++ b/tests/test_tron_tools.py @@ -997,6 +997,17 @@ def test_format_tron_action_dict_paasta(self): "mem": 1200, "disk": 42, "env": mock.ANY, + "topology_spread_constraints": [ + { + "label_selector": { + "app.kubernetes.io/managed-by": "tron", + "paasta.yelp.com/pool": "special_pool", + }, + "max_skew": 1, + "topology_key": "topology.kubernetes.io/zone", + "when_unsatisfiable": "ScheduleAnyway", + }, + ], "secret_volumes": [ { "secret_volume_name": "tron-secret-my--service-secret1", @@ -1347,6 +1358,17 @@ def test_format_tron_action_dict_spark( "KUBECONFIG": "/etc/kubernetes/spark.conf", "AWS_DEFAULT_REGION": "us-west-2", }, + "topology_spread_constraints": [ + { + "label_selector": { + "app.kubernetes.io/managed-by": "tron", + "paasta.yelp.com/pool": "stable", + }, + "max_skew": 1, + "topology_key": "topology.kubernetes.io/zone", + "when_unsatisfiable": "ScheduleAnyway", + }, + ], "node_selectors": {"yelp.com/pool": "stable"}, "cap_add": [], "cap_drop": [ @@ -1474,6 +1496,17 @@ def test_format_tron_action_dict_paasta_k8s_service_account(self): }, "node_selectors": {"yelp.com/pool": "default"}, "env": mock.ANY, + "topology_spread_constraints": [ + { + "label_selector": { + "app.kubernetes.io/managed-by": "tron", + "paasta.yelp.com/pool": "default", + }, + "max_skew": 1, + "topology_key": "topology.kubernetes.io/zone", + "when_unsatisfiable": "ScheduleAnyway", + }, + ], "secret_env": {}, "field_selector_env": {"PAASTA_POD_IP": {"field_path": "status.podIP"}}, "secret_volumes": [], @@ -1621,6 +1654,17 @@ def test_format_tron_action_dict_paasta_k8s( } ], "env": mock.ANY, + "topology_spread_constraints": [ + { + "label_selector": { + "app.kubernetes.io/managed-by": "tron", + "paasta.yelp.com/pool": "special_pool", + }, + "max_skew": 1, + "topology_key": "topology.kubernetes.io/zone", + "when_unsatisfiable": "ScheduleAnyway", + }, + ], "secret_env": { "SOME_SECRET": { "secret_name": "tron-secret-my--service-secret--name", @@ -1719,6 +1763,17 @@ def test_format_tron_action_dict_paasta_no_branch_dict(self): "mem": 1200, "disk": 42, "env": mock.ANY, + "topology_spread_constraints": [ + { + "label_selector": { + "app.kubernetes.io/managed-by": "tron", + "paasta.yelp.com/pool": "special_pool", + }, + "max_skew": 1, + "topology_key": "topology.kubernetes.io/zone", + "when_unsatisfiable": "ScheduleAnyway", + }, + ], "secret_volumes": [ { "secret_volume_name": "tron-secret-my--service-secret1", From 27b20600159db90d1a4c43e653274472200c9bc1 Mon Sep 17 00:00:00 2001 From: Luis Perez Date: Fri, 8 Nov 2024 07:43:55 -0800 Subject: [PATCH 3/5] Wrap Tron TSC code with a feature toggle This'll let us deploy this without big-bouncing Tron everywhere --- paasta_tools/tron_tools.py | 66 +++++++++++++++++++++----------------- paasta_tools/utils.py | 4 +++ 2 files changed, 40 insertions(+), 30 deletions(-) diff --git a/paasta_tools/tron_tools.py b/paasta_tools/tron_tools.py index f5918fce25..dcde970af1 100644 --- a/paasta_tools/tron_tools.py +++ b/paasta_tools/tron_tools.py @@ -610,21 +610,25 @@ def get_node_affinities(self) -> Optional[List[Dict[str, Union[str, List[str]]]] node_selectors = self.config_dict.get("node_selectors", {}) requirements.extend( raw_selectors_to_requirements( - raw_selectors=self.config_dict.get("node_selectors", {}), + raw_selectors=node_selectors, ) ) - # PAASTA-18198: To improve AZ balance with Karpenter, we temporarily allow specifying zone affinities per pool - pool_node_affinities = load_system_paasta_config().get_pool_node_affinities() - if pool_node_affinities and self.get_pool() in pool_node_affinities: - current_pool_node_affinities = pool_node_affinities[self.get_pool()] - # If the service already has a node selector for a zone, we don't want to override it - if current_pool_node_affinities and not contains_zone_label(node_selectors): - requirements.extend( - raw_selectors_to_requirements( - raw_selectors=current_pool_node_affinities, + system_paasta_config = load_system_paasta_config() + if system_paasta_config.get_enable_tron_tsc(): + # PAASTA-18198: To improve AZ balance with Karpenter, we temporarily allow specifying zone affinities per pool + pool_node_affinities = system_paasta_config.get_pool_node_affinities() + if pool_node_affinities and self.get_pool() in pool_node_affinities: + current_pool_node_affinities = pool_node_affinities[self.get_pool()] + # If the service already has a node selector for a zone, we don't want to override it + if current_pool_node_affinities and not contains_zone_label( + node_selectors + ): + requirements.extend( + raw_selectors_to_requirements( + raw_selectors=current_pool_node_affinities, + ) ) - ) if not requirements: return None @@ -983,6 +987,9 @@ def format_tron_action_dict(action_config: TronActionConfig): "service_account_name": action_config.get_service_account_name(), } + # we need this loaded in several branches, so we'll load it once at the start to simplify things + system_paasta_config = load_system_paasta_config() + if executor in KUBERNETES_EXECUTOR_NAMES: # we'd like Tron to be able to distinguish between spark and normal actions # even though they both run on k8s @@ -1004,24 +1011,25 @@ def format_tron_action_dict(action_config: TronActionConfig): result["node_selectors"] = action_config.get_node_selectors() result["node_affinities"] = action_config.get_node_affinities() - # XXX: this is currently hardcoded since we should only really need TSC for zone-aware scheduling - result["topology_spread_constraints"] = [ - { - # try to evenly spread pods across specified topology - "max_skew": 1, - # narrow down what pods to consider when spreading - "label_selector": { - # only consider pods that are managed by tron - "app.kubernetes.io/managed-by": "tron", - # and in the same pool - "paasta.yelp.com/pool": action_config.get_pool(), + if system_paasta_config.get_enable_tron_tsc(): + # XXX: this is currently hardcoded since we should only really need TSC for zone-aware scheduling + result["topology_spread_constraints"] = [ + { + # try to evenly spread pods across specified topology + "max_skew": 1, + # narrow down what pods to consider when spreading + "label_selector": { + # only consider pods that are managed by tron + "app.kubernetes.io/managed-by": "tron", + # and in the same pool + "paasta.yelp.com/pool": action_config.get_pool(), + }, + # now, spread across AZs + "topology_key": "topology.kubernetes.io/zone", + # but if not possible, schedule even with a zonal imbalance + "when_unsatisfiable": "ScheduleAnyway", }, - # now, spread across AZs - "topology_key": "topology.kubernetes.io/zone", - # but if not possible, schedule even with a zonal imbalance - "when_unsatisfiable": "ScheduleAnyway", - }, - ] + ] # XXX: once we're off mesos we can make get_cap_* return just the cap names as a list result["cap_add"] = [cap["value"] for cap in action_config.get_cap_add()] @@ -1068,14 +1076,12 @@ def format_tron_action_dict(action_config: TronActionConfig): # XXX: now that we're actually passing through extra_volumes correctly (e.g., using get_volumes()), # we can get rid of the default_volumes from the Tron master config - system_paasta_config = load_system_paasta_config() extra_volumes = action_config.get_volumes( system_paasta_config.get_volumes(), uses_bulkdata_default=system_paasta_config.get_uses_bulkdata_default(), ) if executor == "spark": is_mrjob = action_config.config_dict.get("mrjob", False) - system_paasta_config = load_system_paasta_config() # inject additional Spark configs in case of Spark commands result["command"] = spark_tools.build_spark_command( result["command"], diff --git a/paasta_tools/utils.py b/paasta_tools/utils.py index 85da5b9777..be1c59aba0 100644 --- a/paasta_tools/utils.py +++ b/paasta_tools/utils.py @@ -2053,6 +2053,7 @@ class SystemPaastaConfigDict(TypedDict, total=False): vitess_throttling_config: Dict uses_bulkdata_default: bool enable_automated_redeploys_default: bool + enable_tron_tsc: bool def load_system_paasta_config( @@ -2821,6 +2822,9 @@ def get_uses_bulkdata_default(self) -> bool: def get_enable_automated_redeploys_default(self) -> bool: return self.config_dict.get("enable_automated_redeploys_default", False) + def get_enable_tron_tsc(self) -> bool: + return self.config_dict.get("enable_tron_tsc", False) + def _run( command: Union[str, List[str]], From 0632a90b78a089048ba086213f5adba897bdb78f Mon Sep 17 00:00:00 2001 From: Luis Perez Date: Fri, 8 Nov 2024 07:47:23 -0800 Subject: [PATCH 4/5] Add feature toggle to tests --- tests/test_tron_tools.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_tron_tools.py b/tests/test_tron_tools.py index 7e9fb54000..31bb358550 100644 --- a/tests/test_tron_tools.py +++ b/tests/test_tron_tools.py @@ -24,6 +24,7 @@ "volumes": [], "dockercfg_location": "/mock/dockercfg", "spark_k8s_role": "spark", + "enable_tron_tsc": True, } ), "/mock/system/configs", From a9274589f4b4993bfa72649edec6e9574c534a11 Mon Sep 17 00:00:00 2001 From: Luis Perez Date: Fri, 8 Nov 2024 08:04:37 -0800 Subject: [PATCH 5/5] Update tests I realized that we weren't consistently using the same mock SystemPaastaConfig - so I went ahead and fixed that while I was here --- tests/test_tron_tools.py | 124 ++++++++++++++++++++++++++++----------- 1 file changed, 90 insertions(+), 34 deletions(-) diff --git a/tests/test_tron_tools.py b/tests/test_tron_tools.py index 31bb358550..11a08270ee 100644 --- a/tests/test_tron_tools.py +++ b/tests/test_tron_tools.py @@ -40,6 +40,7 @@ "tron_k8s_cluster_overrides": { "paasta-dev-test": "paasta-dev", }, + "enable_tron_tsc": True, } ), "/mock/system/configs", @@ -112,7 +113,9 @@ def test_get_env( "paasta_tools.utils.get_service_docker_registry", autospec=True, ), mock.patch( - "paasta_tools.tron_tools.load_system_paasta_config", autospec=True + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, ): env = action_config.get_env() assert not any([env.get("SPARK_OPTS"), env.get("CLUSTERMAN_RESOURCES")]) @@ -288,20 +291,10 @@ def test_get_action_config( "my_job", job_dict, cluster, soa_dir=soa_dir ) - mock_paasta_system_config = utils.SystemPaastaConfig( - config=utils.SystemPaastaConfigDict( - { - "tron_k8s_cluster_overrides": { - "paasta-dev-test": "paasta-dev", - } - } - ), - directory="/mock/system/configs", - ) with mock.patch( "paasta_tools.tron_tools.load_system_paasta_config", autospec=True, - return_value=mock_paasta_system_config, + return_value=MOCK_SYSTEM_PAASTA_CONFIG_OVERRIDES, ): action_config = job_config._get_action_config( "normal", action_dict=action_dict @@ -342,7 +335,11 @@ def test_get_action_config( cluster=expected_cluster, ) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) @mock.patch("paasta_tools.tron_tools.load_v2_deployments_json", autospec=True) def test_get_action_config_load_deployments_false( self, mock_load_deployments, mock_load_system_paasta_config @@ -363,9 +360,6 @@ def test_get_action_config_load_deployments_false( "my_job", job_dict, cluster, load_deployments=False, soa_dir=soa_dir ) mock_load_deployments.side_effect = NoDeploymentsAvailable - mock_load_system_paasta_config.return_value.get_tron_k8s_cluster_overrides.return_value = ( - {} - ) action_config = job_config._get_action_config("normal", action_dict) @@ -531,7 +525,11 @@ def test_format_tron_job_dict_with_cleanup_action( } @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_all_actions( self, mock_load_system_paasta_config, mock_get_pipeline_deploy_groups ): @@ -553,7 +551,11 @@ def test_validate_all_actions( assert len(errors) == 3 @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_invalid_deploy_group( self, mock_load_system_paasta_config, mock_get_pipeline_deploy_groups ): @@ -573,7 +575,11 @@ def test_validate_invalid_deploy_group( assert len(errors) == 1 @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_valid_deploy_group( self, mock_load_system_paasta_config, mock_get_pipeline_deploy_groups ): @@ -593,7 +599,11 @@ def test_validate_valid_deploy_group( assert len(errors) == 0 @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_invalid_action_deploy_group( self, mock_load_system_paasta_config, mock_get_pipeline_deploy_groups ): @@ -617,7 +627,11 @@ def test_validate_invalid_action_deploy_group( assert len(errors) == 1 @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_action_valid_deploy_group( self, mock_load_system_paasta_config, mock_get_pipeline_deploy_groups ): @@ -642,7 +656,11 @@ def test_validate_action_valid_deploy_group( "paasta_tools.tron_tools.TronActionConfig.build_spark_config", autospec=True ) @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_invalid_cpus_in_executor_spark_action( self, mock_load_system_paasta_config, @@ -674,7 +692,11 @@ def test_validate_invalid_cpus_in_executor_spark_action( "paasta_tools.tron_tools.TronActionConfig.build_spark_config", autospec=True ) @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_invalid_mem_in_executor_spark_action( self, mock_load_system_paasta_config, @@ -706,7 +728,11 @@ def test_validate_invalid_mem_in_executor_spark_action( "paasta_tools.tron_tools.TronActionConfig.build_spark_config", autospec=True ) @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_valid_executor_spark_action( self, mock_load_system_paasta_config, @@ -734,7 +760,11 @@ def test_validate_valid_executor_spark_action( assert len(errors) == 0 @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_monitoring( self, mock_load_system_paasta_config, mock_get_pipeline_deploy_groups ): @@ -751,7 +781,11 @@ def test_validate_monitoring( assert len(errors) == 0 @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_monitoring_without_team( self, mock_load_system_paasta_config, mock_get_pipeline_deploy_groups ): @@ -769,7 +803,11 @@ def test_validate_monitoring_without_team( assert job_config.get_monitoring()["team"] == "default_team" @mock.patch("paasta_tools.utils.get_pipeline_deploy_groups", autospec=True) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) def test_validate_monitoring_with_invalid_team( self, mock_load_system_paasta_config, mock_get_pipeline_deploy_groups ): @@ -799,12 +837,15 @@ def test_get_monitoring(self, tronfig_monitoring): class TestTronTools: - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) - def test_load_tron_config(self, mock_system_paasta_config): + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) + def test_load_tron_config(self, mock_load_system_paasta_config): result = tron_tools.load_tron_config() - assert mock_system_paasta_config.return_value.get_tron_config.call_count == 1 assert result == tron_tools.TronConfig( - mock_system_paasta_config.return_value.get_tron_config.return_value + mock_load_system_paasta_config().get_tron_config() ) @mock.patch("paasta_tools.tron_tools.load_tron_config", autospec=True) @@ -909,10 +950,13 @@ def test_format_tron_action_dict_default_executor(self): with mock.patch.object( action_config, "get_docker_registry", return_value="docker-registry.com:400" ), mock.patch( - "paasta_tools.utils.load_system_paasta_config", autospec=True + "paasta_tools.utils.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, ), mock.patch( "paasta_tools.tron_tools.load_system_paasta_config", autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, ), mock.patch( "paasta_tools.tron_tools.add_volumes_for_authenticating_services", autospec=True, @@ -978,6 +1022,7 @@ def test_format_tron_action_dict_paasta(self): ), mock.patch( "paasta_tools.tron_tools.load_system_paasta_config", autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, ), mock.patch( "paasta_tools.tron_tools.add_volumes_for_authenticating_services", autospec=True, @@ -1064,7 +1109,11 @@ def test_format_tron_action_dict_paasta(self): @mock.patch( "paasta_tools.kubernetes_tools.kube_config.load_kube_config", autospec=True ) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) @mock.patch("paasta_tools.tron_tools.get_k8s_url_for_cluster", autospec=True) @mock.patch( "service_configuration_lib.spark_config._get_k8s_docker_volumes_conf", @@ -1466,6 +1515,7 @@ def test_format_tron_action_dict_paasta_k8s_service_account(self): ), mock.patch( "paasta_tools.tron_tools.load_system_paasta_config", autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, ), mock.patch( "paasta_tools.tron_tools.add_volumes_for_authenticating_services", autospec=True, @@ -1610,6 +1660,7 @@ def test_format_tron_action_dict_paasta_k8s( ), mock.patch( "paasta_tools.tron_tools.load_system_paasta_config", autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, ), mock.patch( "paasta_tools.secret_tools.is_shared_secret_from_secret_name", autospec=True, @@ -1748,6 +1799,7 @@ def test_format_tron_action_dict_paasta_no_branch_dict(self): ), mock.patch( "paasta_tools.tron_tools.load_system_paasta_config", autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, ), mock.patch( "paasta_tools.tron_tools.add_volumes_for_authenticating_services", autospec=True, @@ -1853,7 +1905,11 @@ def test_load_tron_service_config_empty(self, mock_read_extra_service_informatio service_name="service", extra_info="tron-test-cluster", soa_dir="fake" ) - @mock.patch("paasta_tools.tron_tools.load_system_paasta_config", autospec=True) + @mock.patch( + "paasta_tools.tron_tools.load_system_paasta_config", + autospec=True, + return_value=MOCK_SYSTEM_PAASTA_CONFIG, + ) @mock.patch("paasta_tools.tron_tools.load_tron_config", autospec=True) @mock.patch("paasta_tools.tron_tools.load_tron_service_config", autospec=True) @mock.patch("paasta_tools.tron_tools.format_tron_job_dict", autospec=True)