From 0f8070a6fdbc5d4ee77253ce58792923d67573c8 Mon Sep 17 00:00:00 2001
From: He Yin Ling <heyinling@espressif.com>
Date: Tue, 20 Nov 2018 23:11:04 +0800
Subject: [PATCH] tiny-test-fw: support capture raw data in DUT:

test cases might want to use `expect` and get raw data from DUT at the same time. New added capture method provides a way to do that.
---
 tools/tiny-test-fw/DUT.py | 53 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 50 insertions(+), 3 deletions(-)

diff --git a/tools/tiny-test-fw/DUT.py b/tools/tiny-test-fw/DUT.py
index e49a121f3..88894fd49 100644
--- a/tools/tiny-test-fw/DUT.py
+++ b/tools/tiny-test-fw/DUT.py
@@ -205,12 +205,14 @@ class _RecvThread(threading.Thread):
 
     PERFORMANCE_PATTERN = re.compile(r"\[Performance]\[(\w+)]: ([^\r\n]+)\r?\n")
 
-    def __init__(self, read, data_cache):
+    def __init__(self, read, data_cache, recorded_data, record_data_lock):
         super(_RecvThread, self).__init__()
         self.exit_event = threading.Event()
         self.setDaemon(True)
         self.read = read
         self.data_cache = data_cache
+        self.recorded_data = recorded_data
+        self.record_data_lock = record_data_lock
         # cache the last line of recv data for collecting performance
         self._line_cache = str()
 
@@ -243,7 +245,10 @@ class _RecvThread(threading.Thread):
         while not self.exit_event.isSet():
             data = self.read(1000)
             if data:
-                self.data_cache.put(data)
+                with self.record_data_lock:
+                    self.data_cache.put(data)
+                    for capture_id in self.recorded_data:
+                        self.recorded_data[capture_id].put(data)
                 self.collect_performance(data)
 
     def exit(self):
@@ -274,6 +279,11 @@ class BaseDUT(object):
         self.log_file = log_file
         self.app = app
         self.data_cache = _DataCache()
+        # the main process of recorded data are done in receive thread
+        # but receive thread could be closed in DUT lifetime (tool methods)
+        # so we keep it in BaseDUT, as their life cycle are same
+        self.recorded_data = dict()
+        self.record_data_lock = threading.RLock()
         self.receive_thread = None
         self.expect_failures = []
         # open and start during init
@@ -389,7 +399,8 @@ class BaseDUT(object):
         :return: None
         """
         self._port_open()
-        self.receive_thread = _RecvThread(self._port_read, self.data_cache)
+        self.receive_thread = _RecvThread(self._port_read, self.data_cache,
+                                          self.recorded_data, self.record_data_lock)
         self.receive_thread.start()
 
     def close(self):
@@ -448,6 +459,42 @@ class BaseDUT(object):
         self.data_cache.flush(size)
         return data
 
+    def start_capture_raw_data(self, capture_id="default"):
+        """
+        Sometime application want to get DUT raw data and use ``expect`` method at the same time.
+        Capture methods provides a way to get raw data without affecting ``expect`` or ``read`` method.
+
+        If you call ``start_capture_raw_data`` with same capture id again, it will restart capture on this ID.
+
+        :param capture_id: ID of capture. You can use different IDs to do different captures at the same time.
+        """
+        with self.record_data_lock:
+            try:
+                # if start capture on existed ID, we do flush data and restart capture
+                self.recorded_data[capture_id].flush()
+            except KeyError:
+                # otherwise, create new data cache
+                self.recorded_data[capture_id] = _DataCache()
+
+    def stop_capture_raw_data(self, capture_id="default"):
+        """
+        Stop capture and get raw data.
+        This method should be used after ``start_capture_raw_data`` on the same capture ID.
+
+        :param capture_id: ID of capture.
+        :return: captured raw data between start capture and stop capture.
+        """
+        with self.record_data_lock:
+            try:
+                ret = self.recorded_data[capture_id].get_data()
+                self.recorded_data.pop(capture_id)
+            except KeyError as e:
+                e.message = "capture_id does not exist. " \
+                            "You should call start_capture_raw_data with same ID " \
+                            "before calling stop_capture_raw_data"
+                raise e
+        return ret
+
     # expect related methods
 
     @staticmethod