From 0d6c6883dd0d3172f8ecc57a1ca6041e9d0043b5 Mon Sep 17 00:00:00 2001
From: Marc Bonnici <marc.bonnici@arm.com>
Date: Thu, 3 Oct 2019 11:21:03 +0100
Subject: [PATCH] instruments: Add `keep_raw` parameter to control teardown
 deletion

Add a `keep_raw` parameter that prevents raw files from being deleted
during teardown in case they are still required.
---
 devlib/instrument/acmecape.py         | 9 ++++++---
 devlib/instrument/arm_energy_probe.py | 8 +++++---
 devlib/instrument/daq.py              | 7 +++++--
 devlib/instrument/energy_probe.py     | 7 +++++--
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/devlib/instrument/acmecape.py b/devlib/instrument/acmecape.py
index e60dda2..4a0a709 100644
--- a/devlib/instrument/acmecape.py
+++ b/devlib/instrument/acmecape.py
@@ -58,12 +58,14 @@ class AcmeCapeInstrument(Instrument):
                  iio_capture=which('iio-capture'),
                  host='baylibre-acme.local',
                  iio_device='iio:device0',
-                 buffer_size=256):
+                 buffer_size=256,
+                 keep_raw=False):
         super(AcmeCapeInstrument, self).__init__(target)
         self.iio_capture = iio_capture
         self.host = host
         self.iio_device = iio_device
         self.buffer_size = buffer_size
+        self.keep_raw = keep_raw
         self.sample_rate_hz = 100
         if self.iio_capture is None:
             raise HostError('Missing iio-capture binary')
@@ -161,5 +163,6 @@ class AcmeCapeInstrument(Instrument):
         return [self.raw_data_file]
 
     def teardown(self):
-        if os.path.isfile(self.raw_data_file):
-            os.remove(self.raw_data_file)
+        if not self.keep_raw:
+            if os.path.isfile(self.raw_data_file):
+                os.remove(self.raw_data_file)
diff --git a/devlib/instrument/arm_energy_probe.py b/devlib/instrument/arm_energy_probe.py
index 093999b..0c57407 100644
--- a/devlib/instrument/arm_energy_probe.py
+++ b/devlib/instrument/arm_energy_probe.py
@@ -71,7 +71,7 @@ class ArmEnergyProbeInstrument(Instrument):
 
     MAX_CHANNELS = 12 # 4 Arm Energy Probes
 
-    def __init__(self, target, config_file='./config-aep', ):
+    def __init__(self, target, config_file='./config-aep', keep_raw=False):
         super(ArmEnergyProbeInstrument, self).__init__(target)
         self.arm_probe = which('arm-probe')
         if self.arm_probe is None:
@@ -80,6 +80,7 @@ class ArmEnergyProbeInstrument(Instrument):
         self.attributes = ['power', 'voltage', 'current']
         self.sample_rate_hz = 10000
         self.config_file = config_file
+        self.keep_raw = keep_raw
 
         self.parser = AepParser()
         #TODO make it generic
@@ -144,5 +145,6 @@ class ArmEnergyProbeInstrument(Instrument):
         return [self.output_file_raw]
 
     def teardown(self):
-        if os.path.isfile(self.output_file_raw):
-            os.remove(self.output_file_raw)
+        if not self.keep_raw:
+            if os.path.isfile(self.output_file_raw):
+                os.remove(self.output_file_raw)
diff --git a/devlib/instrument/daq.py b/devlib/instrument/daq.py
index 68cad46..1a5b7c3 100644
--- a/devlib/instrument/daq.py
+++ b/devlib/instrument/daq.py
@@ -45,9 +45,11 @@ class DaqInstrument(Instrument):
                  dv_range=0.2,
                  sample_rate_hz=10000,
                  channel_map=(0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23),
+                 keep_raw=False
                  ):
         # pylint: disable=no-member
         super(DaqInstrument, self).__init__(target)
+        self.keep_raw = keep_raw
         self._need_reset = True
         self._raw_files = []
         if execute_command is None:
@@ -155,8 +157,9 @@ class DaqInstrument(Instrument):
 
     def teardown(self):
         self.execute('close')
-        if os.path.isdir(tempdir):
-            shutil.rmtree(tempdir)
+        if not self.keep_raw:
+            if os.path.isdir(tempdir):
+                shutil.rmtree(tempdir)
 
     def execute(self, command, **kwargs):
         return execute_command(self.server_config, command, **kwargs)
diff --git a/devlib/instrument/energy_probe.py b/devlib/instrument/energy_probe.py
index 63b0ec6..07fe24b 100644
--- a/devlib/instrument/energy_probe.py
+++ b/devlib/instrument/energy_probe.py
@@ -34,9 +34,11 @@ class EnergyProbeInstrument(Instrument):
     def __init__(self, target, resistor_values,
                  labels=None,
                  device_entry='/dev/ttyACM0',
+                 keep_raw=False
                  ):
         super(EnergyProbeInstrument, self).__init__(target)
         self.resistor_values = resistor_values
+        self.keep_raw = keep_raw
         if labels is not None:
             self.labels = labels
         else:
@@ -128,5 +130,6 @@ class EnergyProbeInstrument(Instrument):
         return [self.raw_data_file]
 
     def teardown(self):
-        if os.path.isfile(self.raw_data_file):
-            os.remove(self.raw_data_file)
+        if self.keep_raw:
+            if os.path.isfile(self.raw_data_file):
+                os.remove(self.raw_data_file)