diff --git a/ec2init/plans/ndb/ndb_init_plans.py b/ec2init/plans/ndb/ndb_init_plans.py
index e439ba5..ef67e91 100644
--- a/ec2init/plans/ndb/ndb_init_plans.py
+++ b/ec2init/plans/ndb/ndb_init_plans.py
@@ -5,48 +5,67 @@ from typing import Sequence, List, Optional
 
 from ec2init import Procedure
 from ec2init.plans.init.base_init_plan import BaseInitPlan
-from ec2init.plans.ndb.start_agent import BaseStartNdbAgent, BootstrapMgmStartNdbAgent, StartNdbAgent
+from ec2init.plans.ndb.start_agent import StartNdbAgent
 from ec2init.plans.init.user_init_script import UserInitScript
 from ec2init.plans.init.callback_hopsworks import CallbackHopsworks
 
 logger = logging.getLogger(__name__)
 
+
 class NdbNodeInitPlan(BaseInitPlan):
     def init(self, name: str, config: configparser.ConfigParser, **kwargs: str) -> None:
         super().init(name, config, **kwargs)
         self.extra_mysqld_vms = kwargs.get("extra_mysqld_vms", "1")
         logger.info("%s - extra_mysqld_vms = %s", self.name, self.extra_mysqld_vms)
-    
+
     def create(self) -> Sequence[Procedure]:
         plan: List[Procedure] = []
 
         self.add_callback_instance_booted(plan, user_email=self.user_email)
         callback_hopsworks_check = self.plan_utils.get_procedure_check_file(
-            "{}_booted".format(CallbackHopsworks.__name__))
+            "{}_booted".format(CallbackHopsworks.__name__)
+        )
         if self.plan_utils.check_file_exists(callback_hopsworks_check):
             self.add_user_init_script(plan)
         return plan
 
     def add_start_ndb_agent(self, plan: List[Procedure]) -> None:
-        user_init_check = self.plan_utils.get_procedure_check_file(UserInitScript.__name__)
-        if self.plan_utils.check_file_exists(user_init_check) and not BaseStartNdbAgent.has_run():
+        user_init_check = self.plan_utils.get_procedure_check_file(
+            UserInitScript.__name__
+        )
+        if (
+            self.plan_utils.check_file_exists(user_init_check)
+            and not StartNdbAgent.has_run()
+        ):
             logger.info(f"{self.name} - Adding procedure {StartNdbAgent.__name__}")
-            sna = StartNdbAgent(StartNdbAgent.__name__, self.config, self.api_key, self.extra_mysqld_vms)
+            sna = StartNdbAgent(
+                StartNdbAgent.__name__,
+                self.config,
+                self.api_key,
+                False,
+                self.gateway,
+                self.uid,
+                self.instance_id,
+                self.extra_mysqld_vms,
+            )
             sna.friendly_name = "Setting NDB agent"
             sna.report_callback = self.plan_utils.create_callback_procedure(
-                "ndb_starting", stage="Setting NDB agent", instance_type=self.get_instance_type())
+                "ndb_starting",
+                stage="Setting NDB agent",
+                instance_type=self.get_instance_type(),
+            )
             plan.append(sna)
 
 
-
 class MgmNodeInitPlan(NdbNodeInitPlan):
     def init(self, name: str, config: configparser.ConfigParser, **kwargs: str) -> None:
         super().init(name, config, **kwargs)
         self.instance_type = "ndb_mgm"
-        self.is_bootstrap = True if kwargs.get("mgm_bootstrap", "false").lower() == "true" else False
+        self.is_bootstrap = (
+            True if kwargs.get("mgm_bootstrap", "false").lower() == "true" else False
+        )
         if self.is_bootstrap:
             self.instance_type += "/bootstrap"
-        
 
     def create(self) -> Sequence[Procedure]:
         plan: List[Procedure] = []
@@ -56,21 +75,29 @@ class MgmNodeInitPlan(NdbNodeInitPlan):
 
         self.add_start_ndb_agent(plan)
         return plan
-    
+
     def add_start_ndb_agent(self, plan: List[Procedure]) -> None:
-        if not BaseStartNdbAgent.has_run():
+        if not StartNdbAgent.has_run():
             logger.info(f"{self.name} - Adding procedure {StartNdbAgent.__name__}")
-            sna: Optional[BaseStartNdbAgent] = None
-            if self.is_bootstrap:
-                sna = BootstrapMgmStartNdbAgent(BootstrapMgmStartNdbAgent.__name__, self.config,
-                    self.api_key, self.gateway, self.uid, self.instance_id, self.extra_mysqld_vms)
-            else:
-                sna = StartNdbAgent(StartNdbAgent.__name__, self.config, self.api_key, self.extra_mysqld_vms)
+            sna = StartNdbAgent(
+                StartNdbAgent.__name__,
+                self.config,
+                self.api_key,
+                self.is_bootstrap,
+                self.gateway,
+                self.uid,
+                self.instance_id,
+                self.extra_mysqld_vms,
+            )
             sna.friendly_name = "Setting NDB agent"
             sna.report_callback = self.plan_utils.create_callback_procedure(
-                "ndb_starting", stage="Setting NDB agent", instance_type=self.get_instance_type())
+                "ndb_starting",
+                stage="Setting NDB agent",
+                instance_type=self.get_instance_type(),
+            )
             plan.append(sna)
 
+
 class NdbdNodeInitPlan(NdbNodeInitPlan):
     def init(self, name: str, config: configparser.ConfigParser, **kwargs: str) -> None:
         super().init(name, config, **kwargs)
@@ -85,6 +112,7 @@ class NdbdNodeInitPlan(NdbNodeInitPlan):
 
         return plan
 
+
 class MysqldNodeInitPlan(NdbNodeInitPlan):
     def init(self, name: str, config: configparser.ConfigParser, **kwargs: str) -> None:
         super().init(name, config, **kwargs)
@@ -98,6 +126,7 @@ class MysqldNodeInitPlan(NdbNodeInitPlan):
         self.add_start_ndb_agent(plan)
         return plan
 
+
 class ApiNodeInitPlan(NdbNodeInitPlan):
     def init(self, name: str, config: configparser.ConfigParser, **kwargs: str) -> None:
         super().init(name, config, **kwargs)
@@ -111,8 +140,11 @@ class ApiNodeInitPlan(NdbNodeInitPlan):
         self.add_start_ndb_agent(plan)
         return plan
 
+
 class RonDBAllInOneInitPlan(MgmNodeInitPlan):
     def init(self, name: str, config: configparser.ConfigParser, **kwargs: str) -> None:
         super().init(name, config, **kwargs)
         self.instance_type = "rondb_all-in-one"
-        self.is_bootstrap = True if kwargs.get("mgm_bootstrap", "false").lower() == "true" else False
\ No newline at end of file
+        self.is_bootstrap = (
+            True if kwargs.get("mgm_bootstrap", "false").lower() == "true" else False
+        )
diff --git a/ec2init/plans/ndb/start_agent.py b/ec2init/plans/ndb/start_agent.py
index c708329..bf27dd0 100644
--- a/ec2init/plans/ndb/start_agent.py
+++ b/ec2init/plans/ndb/start_agent.py
@@ -12,44 +12,60 @@ logger = logging.getLogger(__name__)
 
 HEARTBEAT_ENDPOINT = "{gateway}/eventsv2/ndbagent/{uid}/{cid}/heartbeat"
 
-class BaseStartNdbAgent(Procedure):
-    AGENT_CONFIG="/srv/hops/ndb-agent/ndb-agent/config/config.yml"
-    def __init__(self, name: str, config: configparser.ConfigParser, api_key: str, hb_endpoint: Optional[str], extra_mysqld_vms: str) -> None:
+
+class StartNdbAgent(Procedure):
+    AGENT_CONFIG = "/srv/hops/ndb-agent/ndb-agent/config/config.yml"
+
+    def __init__(
+        self,
+        name: str,
+        config: configparser.ConfigParser,
+        api_key: str,
+        is_bootstrap: bool,
+        gateway: str,
+        uid: str,
+        cluster_id: str,
+        extra_mysqld_vms: str,
+    ) -> None:
         self.name = name
         self.config = config
         self.api_key = api_key
-        self.hb_endpoint = hb_endpoint
+        self.is_bootstrap = is_bootstrap
+        self.hb_endpoint = HEARTBEAT_ENDPOINT.format(
+            gateway=gateway, uid=uid, cid=cluster_id
+        )
         self.extra_mysqld_vms = extra_mysqld_vms
 
-
     def execute(self) -> List[Procedure]:
         logger.debug(f"{self.name} - Replacing API key")
         subString = f"s/API_KEY/{self.api_key}/g"
-        cmd = Seds.REPLACE.format(subString, BaseStartNdbAgent.AGENT_CONFIG)
+        cmd = Seds.REPLACE.format(subString, StartNdbAgent.AGENT_CONFIG)
         SystemUtils.execute_command(cmd.split(" "))
         logger.info(f"{self.name} - Replaced API key")
 
         logger.debug(f"{self.name} - Configure extra_mysqld_vms")
-        subString = f"s/'extra-mysqld-vms:.*'/'extra-mysqld-vms: {self.extra_mysqld_vms}'/g"
-        cmd = Seds.REPLACE.format(subString, BaseStartNdbAgent.AGENT_CONFIG)
+        subString = (
+            f"s/'extra-mysqld-vms:.*'/'extra-mysqld-vms: {self.extra_mysqld_vms}'/g"
+        )
+        cmd = Seds.REPLACE.format(subString, StartNdbAgent.AGENT_CONFIG)
         SystemUtils.execute_command([cmd], shell=True)
 
-        if self.hb_endpoint:
+        logger.debug(f"{self.name} - Replacing Heartbeat endpoint")
+        subString = f"s;HEARTBEAT_ENDPOINT;{str(self.hb_endpoint)};g"
+        cmd = Seds.REPLACE.format(subString, StartNdbAgent.AGENT_CONFIG)
+        SystemUtils.execute_command(cmd.split(" "))
+        logger.info(f"{self.name} - Replaced Heartbeat endpoint")
+        
+        if self.is_bootstrap:
             logger.debug(f"{self.name} - Enabling Heartbeat")
             subString = f"s/HEARTBEAT_ENABLED/true/g"
-            cmd = Seds.REPLACE.format(subString, BaseStartNdbAgent.AGENT_CONFIG)
+            cmd = Seds.REPLACE.format(subString, StartNdbAgent.AGENT_CONFIG)
             SystemUtils.execute_command(cmd.split(" "))
             logger.info(f"{self.name} - Enabled Heartbeat")
-
-            logger.debug(f"{self.name} - Replacing Heartbeat endpoint")
-            subString = f"s;HEARTBEAT_ENDPOINT;{str(self.hb_endpoint)};g"
-            cmd = Seds.REPLACE.format(subString, BaseStartNdbAgent.AGENT_CONFIG)
-            SystemUtils.execute_command(cmd.split(" "))
-            logger.info(f"{self.name} - Replaced Heartbeat endpoint")
         else:
             logger.debug(f"{self.name} - Disabling Heartbeat")
             subString = f"s/HEARTBEAT_ENABLED/false/g"
-            cmd = Seds.REPLACE.format(subString, BaseStartNdbAgent.AGENT_CONFIG)
+            cmd = Seds.REPLACE.format(subString, StartNdbAgent.AGENT_CONFIG)
             SystemUtils.execute_command(cmd.split(" "))
             logger.info(f"{self.name} - Disabled Heartbeat")
 
@@ -65,14 +81,6 @@ class BaseStartNdbAgent(Procedure):
 
     @staticmethod
     def has_run() -> bool:
-        return StartServices._is_service_running("ndb-agent")\
-            and StartServices.is_service_enabled("ndb-agent")
-
-class BootstrapMgmStartNdbAgent(BaseStartNdbAgent):
-    def __init__(self, name: str, config: configparser.ConfigParser, api_key: str, gateway: str, uid: str, cluster_id: str, extra_mysqld_vms: str) -> None:
-        hb_endpoint = HEARTBEAT_ENDPOINT.format(gateway=gateway, uid=uid, cid=cluster_id)
-        super().__init__(name, config, api_key, hb_endpoint, extra_mysqld_vms)
-
-class StartNdbAgent(BaseStartNdbAgent):
-    def __init__(self, name: str, config: configparser.ConfigParser, api_key: str, extra_mysqld_vms: str) -> None:
-        super().__init__(name, config, api_key, None, extra_mysqld_vms)
\ No newline at end of file
+        return StartServices._is_service_running(
+            "ndb-agent"
+        ) and StartServices.is_service_enabled("ndb-agent")
