diff --git a/lib/ansible/cli/scripts/ansible_connection_cli_stub.py b/lib/ansible/cli/scripts/ansible_connection_cli_stub.py
index 2505cbf2c46..b5347fd9f3f 100755
--- a/lib/ansible/cli/scripts/ansible_connection_cli_stub.py
+++ b/lib/ansible/cli/scripts/ansible_connection_cli_stub.py
@@ -306,9 +306,9 @@ def main():
                 pc_data = to_text(init_data)
                 try:
                     conn.update_play_context(pc_data)
-                    conn.update_cli_prompt_context(task_uuid)
+                    conn.set_check_prompt(task_uuid)
                 except Exception as exc:
-                    # Only network_cli has update_play context and set_cli_prompt_context, so missing this is
+                    # Only network_cli has update_play context and set_check_prompt, so missing this is
                     # not fatal e.g. netconf
                     if isinstance(exc, ConnectionError) and getattr(exc, 'code', None) == -32601:
                         pass
diff --git a/lib/ansible/plugins/connection/network_cli.py b/lib/ansible/plugins/connection/network_cli.py
index 4d36c7d82b8..199da9988f7 100644
--- a/lib/ansible/plugins/connection/network_cli.py
+++ b/lib/ansible/plugins/connection/network_cli.py
@@ -273,6 +273,7 @@ options:
         - name: ansible_network_cli_retries
 """
 
+from functools import wraps
 import getpass
 import json
 import logging
@@ -290,10 +291,20 @@ from ansible.module_utils.six.moves import cPickle
 from ansible.module_utils.network.common.utils import to_list
 from ansible.module_utils._text import to_bytes, to_text
 from ansible.playbook.play_context import PlayContext
-from ansible.plugins.connection import NetworkConnectionBase, ensure_connect
+from ansible.plugins.connection import NetworkConnectionBase
 from ansible.plugins.loader import cliconf_loader, terminal_loader, connection_loader
 
 
+def ensure_connect(func):
+    @wraps(func)
+    def wrapped(self, *args, **kwargs):
+        if not self._connected:
+            self._connect()
+        self.update_cli_prompt_context()
+        return func(self, *args, **kwargs)
+    return wrapped
+
+
 class AnsibleCmdRespRecv(Exception):
     pass
 
@@ -319,6 +330,9 @@ class Connection(NetworkConnectionBase):
         self._terminal = None
         self.cliconf = None
         self._paramiko_conn = None
+
+        # Managing prompt context
+        self._check_prompt = False
         self._task_uuid = to_text(kwargs.get('task_uuid', ''))
 
         if self._play_context.verbosity > 3:
@@ -409,11 +423,14 @@ class Connection(NetworkConnectionBase):
         if hasattr(self, 'disable_response_logging'):
             self.disable_response_logging()
 
-    def update_cli_prompt_context(self, task_uuid):
+    def set_check_prompt(self, task_uuid):
+        self._check_prompt = task_uuid
+
+    def update_cli_prompt_context(self):
         # set cli prompt context at the start of new task run only
-        if self._task_uuid != task_uuid:
+        if self._check_prompt and self._task_uuid != self._check_prompt:
+            self._task_uuid, self._check_prompt = self._check_prompt, False
             self.set_cli_prompt_context()
-            self._task_uuid = task_uuid
 
     def _connect(self):
         '''
diff --git a/test/integration/targets/eos_smoke/tests/cli/reboot.yaml b/test/integration/targets/eos_smoke/tests/cli/reboot.yaml
new file mode 100644
index 00000000000..a83c7815d4a
--- /dev/null
+++ b/test/integration/targets/eos_smoke/tests/cli/reboot.yaml
@@ -0,0 +1,20 @@
+---
+- block:
+  - cli_command:
+      command: reload power
+      prompt:
+        - "yes/no/cancel/diff]"
+        - "confirm]"
+      answer:
+        - "no"
+        - ""
+      check_all: yes
+    become: yes
+
+  - wait_for_connection:
+      delay: 20
+      sleep: 10
+
+  - cli_command:
+      command: show version
+  when: ansible_connection.endswith("network_cli")
diff --git a/test/integration/targets/ios_smoke/tests/cli/reboot.yaml b/test/integration/targets/ios_smoke/tests/cli/reboot.yaml
new file mode 100644
index 00000000000..6f9a77b255c
--- /dev/null
+++ b/test/integration/targets/ios_smoke/tests/cli/reboot.yaml
@@ -0,0 +1,19 @@
+---
+- block:
+  - cli_command:
+      command: reload
+      prompt:
+        - "yes/no"
+        - "confirm"
+      answer:
+        - "no"
+        - "y"
+      check_all: yes
+
+  - wait_for_connection:
+      delay: 20
+      sleep: 10
+
+  - cli_command:
+      command: show version
+  when: ansible_connection.endswith("network_cli")
diff --git a/test/integration/targets/junos_smoke/tasks/cli.yaml b/test/integration/targets/junos_smoke/tasks/cli.yaml
new file mode 100644
index 00000000000..7b0afac3acf
--- /dev/null
+++ b/test/integration/targets/junos_smoke/tasks/cli.yaml
@@ -0,0 +1,15 @@
+- name: collect cli test cases
+  find:
+    paths: "{{ role_path }}/tests/cli"
+    patterns: "{{ testcase }}.yaml"
+  connection: local
+  register: test_cases
+
+- name: set test_items
+  set_fact: test_items="{{ test_cases.files | map(attribute='path') | list }}"
+
+- name: run test case (connection=network_cli)
+  include: "{{ test_case_to_run }} ansible_connection=network_cli"
+  with_items: "{{ test_items }}"
+  loop_control:
+    loop_var: test_case_to_run
diff --git a/test/integration/targets/junos_smoke/tasks/main.yaml b/test/integration/targets/junos_smoke/tasks/main.yaml
index cc27f174fd8..af08869c922 100644
--- a/test/integration/targets/junos_smoke/tasks/main.yaml
+++ b/test/integration/targets/junos_smoke/tasks/main.yaml
@@ -1,2 +1,3 @@
 ---
+- { include: cli.yaml, tags: ['cli'] }
 - { include: netconf.yaml, tags: ['netconf'] }
diff --git a/test/integration/targets/junos_smoke/tests/cli/reboot.yaml b/test/integration/targets/junos_smoke/tests/cli/reboot.yaml
new file mode 100644
index 00000000000..77c41f3b243
--- /dev/null
+++ b/test/integration/targets/junos_smoke/tests/cli/reboot.yaml
@@ -0,0 +1,14 @@
+---
+- cli_command:
+    command: request system reboot
+    prompt:
+      - Reboot the system?
+    answer:
+     - y
+
+- wait_for_connection:
+    delay: 20
+    sleep: 10
+
+- cli_command:
+    command: show version