diff --git a/playbooks/demo_reboot.yml b/playbooks/demo_reboot.yml
index 7734728b..adf58d48 100644
--- a/playbooks/demo_reboot.yml
+++ b/playbooks/demo_reboot.yml
@@ -8,22 +8,16 @@
     post_reboot_delay_v: 0
     reboot_timeout_v: 0
     test_command_v: ""
-    original: 1
   collections:
     - ibm.power_aix
   tasks:
       
-    # - name: "reboot the machine"
-    #   reboot:
-    #           pre_reboot_delay: "{{ pre_reboot_delay_v }}"
-    #           post_reboot_delay: "{{ post_reboot_delay_v }}"
-    #           reboot_timeout: "{{ reboot_timeout_v }}"
-    #           test_command: "{{ test_command_v }}"
-    #   register: reboot_result
+    - name: "reboot the machine"
+      reboot:
+              pre_reboot_delay: "{{ pre_reboot_delay_v }}"
+              post_reboot_delay: "{{ post_reboot_delay_v }}"
+              reboot_timeout: "{{ reboot_timeout_v }}"
+              test_command: "{{ test_command_v }}"
+      register: reboot_result
 
-    # - debug: var=reboot_result
-
-    - name: "Print statement to not use this playbook"
-      ansible.builtin.fail:
-        msg: "Please use demo_reboot_workaround.yml"
-      when: "{{ original }}"
\ No newline at end of file
+    - debug: var=reboot_result
diff --git a/playbooks/demo_reboot_workaround.yml b/playbooks/demo_reboot_workaround.yml
deleted file mode 100644
index 3e4ead41..00000000
--- a/playbooks/demo_reboot_workaround.yml
+++ /dev/null
@@ -1,41 +0,0 @@
----
-- name: "Reboot the machine"
-  hosts: aix
-  gather_facts: no
-  vars:
-    log_file: "/tmp/ansible_reboot_debug.log"
-    pre_reboot_delay_v: 0
-    post_reboot_delay_v: 0
-    reboot_timeout_v: 0
-    test_command_v: ""
-    sleep_time: 5 # This is the time that system will take for starting up.
-  collections:
-    - ibm.power_aix
-  tasks:
-      
-    - name: "reboot the machine"
-      ibm.power_aix.reboot:
-              pre_reboot_delay: "{{ pre_reboot_delay_v }}"
-              post_reboot_delay: "{{ post_reboot_delay_v }}"
-              reboot_timeout: "{{ reboot_timeout_v }}"
-              test_command: "{{ test_command_v }}"
-      register: reboot_result
-
-    - debug: var=reboot_result
-
-    - name: "Sleep for {{sleep_time}} minutes"
-      ansible.builtin.pause:
-        minutes: "{{sleep_time}}"
-
-    - name: "Validate reboot"
-      ansible.builtin.raw: "whoami"
-      register: validate_result
-      changed_when: false
-
-    - name: "Fail if reboot could not be validated"
-      ansible.builtin.fail:
-        msg: "Reboot could not be validated."
-      when: validate_result.rc != 0
-
-    - ansible.builtin.debug:
-        msg: "Reboot successful, machine is up now !"
diff --git a/plugins/action/reboot.py b/plugins/action/reboot.py
index 41f74440..3ab4d733 100644
--- a/plugins/action/reboot.py
+++ b/plugins/action/reboot.py
@@ -69,51 +69,51 @@ def perform_reboot(self, pre_reboot_delay):
 
         return reboot_result
 
-    # def validate_reboot(self, test_command, reboot_timeout=None, action_kwargs=None):
-    #     self._display.vvv('Validating reboot....')
-    #     result = {}
-
-    #     try:
-    #         if reboot_timeout is None:
-    #             reboot_timeout = 300
-
-    #         self.do_until_success_or_timeout(
-    #             action_desc="post-reboot test command",
-    #             reboot_timeout=reboot_timeout,
-    #             test_command=test_command,
-    #             action_kwargs=action_kwargs)
-
-    #         result['rebooted'] = True
-    #         result['changed'] = True
-    #         result['msg'] = "System has been rebooted SUCCESSFULLY"
-
-    #     except TimedOutException:
-    #         result['failed'] = True
-    #         result['rebooted'] = True
-    #         result['msg'] = "Reboot Validation failed due to timeout"
-    #         return result
-
-    #     return result
-
-    # def do_until_success_or_timeout(self, action_desc, reboot_timeout, test_command, action_kwargs=None):
-    #     max_end_time = datetime.utcnow() + timedelta(seconds=reboot_timeout)
-    #     if action_kwargs is None:
-    #         action_kwargs = {}
-    #     if test_command is None:
-    #         test_command = 'whoami'
-
-    #     while datetime.utcnow() < max_end_time:
-    #         try:
-    #             result = self._low_level_execute_command(test_command, sudoable=True)
-    #             if result['rc'] == 0:
-    #                 return
-    #         except Exception as e:
-    #             if isinstance(e, AnsibleConnectionFailure):
-    #                 try:
-    #                     self._connection.reset()
-    #                 except AttributeError:
-    #                     pass
-    #     raise TimedOutException("Connection reset failed while validating the reboot.")
+    def validate_reboot(self, test_command, reboot_timeout=None, action_kwargs=None):
+        self._display.vvv('Validating reboot....')
+        result = {}
+
+        try:
+            if reboot_timeout is None:
+                reboot_timeout = 300
+
+            self.do_until_success_or_timeout(
+                action_desc="post-reboot test command",
+                reboot_timeout=reboot_timeout,
+                test_command=test_command,
+                action_kwargs=action_kwargs)
+
+            result['rebooted'] = True
+            result['changed'] = True
+            result['msg'] = "System has been rebooted SUCCESSFULLY"
+
+        except TimedOutException:
+            result['failed'] = True
+            result['rebooted'] = True
+            result['msg'] = "Reboot Validation failed due to timeout"
+            return result
+
+        return result
+
+    def do_until_success_or_timeout(self, action_desc, reboot_timeout, test_command, action_kwargs=None):
+        max_end_time = datetime.utcnow() + timedelta(seconds=reboot_timeout)
+        if action_kwargs is None:
+            action_kwargs = {}
+        if test_command is None:
+            test_command = 'whoami'
+
+        while datetime.utcnow() < max_end_time:
+            try:
+                result = self._low_level_execute_command(test_command, sudoable=True)
+                if result['rc'] == 0:
+                    return
+            except Exception as e:
+                if isinstance(e, AnsibleConnectionFailure):
+                    try:
+                        self._connection.reset()
+                    except AttributeError:
+                        pass
+        raise TimedOutException("Connection reset failed while validating the reboot.")
 
     def run(self, tmp=None, task_vars=None):
         self._supports_async = True
@@ -146,13 +146,13 @@ def run(self, tmp=None, task_vars=None):
             result['elapsed'] = str(elapsed.seconds) + ' sec'
             return result
 
-        # if self.post_reboot_delay != 0:
-        #     self._display.vvv("waiting for post reboot delay of {delay} seconds".format(delay=self.post_reboot_delay))
-        #     time.sleep(self.post_reboot_delay)
+        if self.post_reboot_delay != 0:
+            self._display.vvv("waiting for post reboot delay of {delay} seconds".format(delay=self.post_reboot_delay))
+            time.sleep(self.post_reboot_delay)
 
-        # result = self.validate_reboot(test_command, reboot_timeout, action_kwargs=None)
+        result = self.validate_reboot(test_command, reboot_timeout, action_kwargs=None)
 
-        # elapsed = datetime.utcnow() - start
-        # result['elapsed'] = str(elapsed.seconds) + ' sec'
+        elapsed = datetime.utcnow() - start
+        result['elapsed'] = str(elapsed.seconds) + ' sec'
 
         return result