From 745dc9499af59702cefc23978806f278bba990ad Mon Sep 17 00:00:00 2001
From: Marc Bonnici <marc.bonnici@arm.com>
Date: Thu, 28 Nov 2019 16:39:19 +0000
Subject: [PATCH] modules/flash: Add a `connect` parameter to the flash method

Adds a `connect` parameter to the flash method to specifiy whether
devlib should attempt to connect to the target after flashing has
completed.
---
 devlib/module/__init__.py | 2 +-
 devlib/module/android.py  | 5 +++--
 devlib/module/vexpress.py | 5 +++--
 doc/modules.rst           | 3 ++-
 4 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/devlib/module/__init__.py b/devlib/module/__init__.py
index b4ab0d1..0ea7119 100644
--- a/devlib/module/__init__.py
+++ b/devlib/module/__init__.py
@@ -91,7 +91,7 @@ class FlashModule(Module):
 
     kind = 'flash'
 
-    def __call__(self, image_bundle=None, images=None, boot_config=None):
+    def __call__(self, image_bundle=None, images=None, boot_config=None, connect=True):
         raise NotImplementedError()
 
 
diff --git a/devlib/module/android.py b/devlib/module/android.py
index bc91077..c0e1bd5 100644
--- a/devlib/module/android.py
+++ b/devlib/module/android.py
@@ -54,7 +54,7 @@ class FastbootFlashModule(FlashModule):
     def probe(target):
         return target.os == 'android'
 
-    def __call__(self, image_bundle=None, images=None, bootargs=None):
+    def __call__(self, image_bundle=None, images=None, bootargs=None, connect=True):
         if bootargs:
             raise ValueError('{} does not support boot configuration'.format(self.name))
         self.prelude_done = False
@@ -67,7 +67,8 @@ class FastbootFlashModule(FlashModule):
             self.logger.debug('flashing {}'.format(partition))
             self._flash_image(self.target, partition, expand_path(image_path))
         fastboot_command('reboot')
-        self.target.connect(timeout=180)
+        if connect:
+            self.target.connect(timeout=180)
 
     def _validate_image_bundle(self, image_bundle):
         if not tarfile.is_tarfile(image_bundle):
diff --git a/devlib/module/vexpress.py b/devlib/module/vexpress.py
index 0f8ab97..6289d26 100644
--- a/devlib/module/vexpress.py
+++ b/devlib/module/vexpress.py
@@ -325,7 +325,7 @@ class VersatileExpressFlashModule(FlashModule):
         self.timeout = timeout
         self.short_delay = short_delay
 
-    def __call__(self, image_bundle=None, images=None, bootargs=None):
+    def __call__(self, image_bundle=None, images=None, bootargs=None, connect=True):
         self.target.hard_reset()
         with open_serial_connection(port=self.target.platform.serial_port,
                                     baudrate=self.target.platform.baudrate,
@@ -346,7 +346,8 @@ class VersatileExpressFlashModule(FlashModule):
             msg = 'Could not deploy images to {}; got: {}'
             raise TargetStableError(msg.format(self.vemsd_mount, e))
         self.target.boot()
-        self.target.connect(timeout=30)
+        if connect:
+            self.target.connect(timeout=30)
 
     def _deploy_image_bundle(self, bundle):
         self.logger.debug('Validating {}'.format(bundle))
diff --git a/doc/modules.rst b/doc/modules.rst
index 264edca..d67d3fc 100644
--- a/doc/modules.rst
+++ b/doc/modules.rst
@@ -322,7 +322,7 @@ FlashModule
 
     "flash"
 
-.. method:: __call__(image_bundle=None, images=None, boot_config=None)
+.. method:: __call__(image_bundle=None, images=None, boot_config=None, connect=True)
 
     Must be implemented by derived classes.
 
@@ -338,6 +338,7 @@ FlashModule
     :param boot_config: Some platforms require specifying boot arguments at the
                         time of flashing the images, rather than during each
                         reboot. For other platforms, this will be ignored.
+    :connect: Specifiy whether to try and connect to the target after flashing.
 
 
 Module Registration