Merge pull request #339 from espressif/feat/improve-expecting-functions

Feat/improve expecting functions

Author: hfudev

docs/usages/expecting.rst

@@ -20,6 +20,10 @@ All of these functions share these possible keyword arguments:
 
    Will raise an exception if the pattern is found in the output, if specified. (Default: None)
 
+-  ``return_what_before_match``
+
+      Will return the bytes read before the match if specified. (Default: False)
+
 *****************************************
  :func:`~pytest_embedded.dut.Dut.expect`
 *****************************************
@@ -130,6 +134,40 @@ What's more, argument ``pattern`` could be a list of all supported types.
 
 If you set ``expect_all`` to ``True``, the :func:`~pytest_embedded.dut.Dut.expect` function would return with a list of returned values of each item.
 
+You can also set ``return_what_before_match`` to ``True`` to get the bytes read before the match, instead of the match object.
+
+.. code:: python
+
+   import pexpect
+
+   def test_expect_before_match(dut):
+       dut.write('this would be redirected')
+
+       res = dut.expect('would', return_what_before_match=True)
+       assert res == b'this '
+
+       res = dut.expect_exact('be ', return_what_before_match=True)
+       assert res == b' '
+
+       res = dut.expect('ected', return_what_before_match=True)
+       assert res == b'redir'
+
+.. hint::
+
+   For better performance when retrieving text before a pattern, use:
+
+   .. code:: python
+
+      before_str = dut.expect('pattern', return_what_before_match=True).decode('utf-8')
+
+   Instead of:
+
+   .. code:: python
+
+      before_str = dut.expect('(.+)pattern').group(1).decode('utf-8')
+
+   The latter performs unnecessary recursive matching of preceding bytes.
+
 ***********************************************
  :func:`~pytest_embedded.dut.Dut.expect_exact`
 ***********************************************

pytest-embedded/pytest_embedded/dut.py

@@ -67,8 +67,17 @@ def write(self, s: AnyStr) -> None:
     def _pexpect_func(func) -> Callable[..., Union[Match, AnyStr]]:
         @functools.wraps(func)
         def wrapper(
-            self, pattern, *args, expect_all: bool = False, not_matching: List[Union[str, re.Pattern]] = (), **kwargs
+            self,
+            pattern,
+            *args,
+            expect_all: bool = False,
+            not_matching: List[Union[str, re.Pattern]] = (),
+            return_what_before_match: bool = False,
+            **kwargs,
         ) -> Union[Union[Match, AnyStr], List[Union[Match, AnyStr]]]:
+            if return_what_before_match and expect_all:
+                raise ValueError('`return_what_before_match` and `expect_all` cannot be `True` at the same time.')
+
             patterns = to_list(pattern)
             res = []
             while patterns:
@@ -101,6 +110,9 @@ def wrapper(
                 else:
                     break  # one succeeded. leave the loop
 
+            if return_what_before_match:
+                return self.pexpect_proc.before
+
             if len(res) == 1:
                 return res[0]
 
@@ -121,6 +133,8 @@ def expect(self, pattern, **kwargs) -> Match:
             expect_all (bool): need to match all specified patterns if this flag is `True`.
                 Otherwise match any of them could pass
             not_matching: string, or compiled regex, or a list of string and compiled regex.
+            return_what_before_match (bool): return the bytes before the matched pattern.
+                Cannot be specified together with `expect_all`
 
         Returns:
             `AnyStr` or `re.Match`
@@ -143,6 +157,8 @@ def expect_exact(self, pattern, **kwargs) -> Match:
             expect_all (bool): need to match all specified patterns if this flag is `True`.
                 Otherwise match any of them could pass
             not_matching: string, or compiled regex, or a list of string and compiled regex.
+            return_what_before_match (bool): return the bytes before the matched pattern.
+                Cannot be specified together with `expect_all`
 
         Returns:
             `AnyStr` or `re.Match`

pytest-embedded/pytest_embedded/dut_factory.py

@@ -50,45 +50,40 @@ def msg_queue_gn() -> MessageQueue:
 
 
 def _listen(q: MessageQueue, filepath: str, with_timestamp: bool = True, count: int = 1, total: int = 1) -> None:
-    _added_prefix = False
+    shall_add_prefix = True
     while True:
-        msgs = q.get_all()
-        if not msgs:
+        msg = q.get()
+        if not msg:
             continue
 
-        msg_b = b''.join(msgs)
         with open(filepath, 'ab') as fw:
-            fw.write(msg_b)
+            fw.write(msg)
             fw.flush()
 
-        _s = to_str(msg_b)
+        _s = to_str(msg)
         if not _s:
             continue
 
         prefix = ''
         if total > 1:
-            source = f'dut-{count}'
-        else:
-            source = None
-
-        if source:
-            prefix = f'[{source}] ' + prefix
+            prefix = f'[dut-{count}] '
 
         if with_timestamp:
             prefix = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') + ' ' + prefix
 
-        if not _added_prefix:
+        if shall_add_prefix:
             _s = prefix + _s
-            _added_prefix = True
+
         _s = _s.replace('\r\n', '\n')  # remove extra \r. since multi-dut \r would mess up the log
-        _s = _s.replace('\n', '\n' + prefix)
-        if prefix and _s.endswith(prefix):
-            _s = _s.rsplit(prefix, maxsplit=1)[0]
-            _added_prefix = False
+        if _s.endswith('\n'):  # complete line
+            shall_add_prefix = True
+            _s = _s[:-1].replace('\n', '\n' + prefix) + '\n'
+        else:
+            shall_add_prefix = False
+            _s = _s.replace('\n', '\n' + prefix)
 
         _stdout.write(_s)
         _stdout.flush()
-        time.sleep(0.05)
 
 
 def _listener_gn(msg_queue, _pexpect_logfile, with_timestamp, dut_index, dut_total) -> multiprocessing.Process:

pytest-embedded/pytest_embedded/log.py

@@ -8,7 +8,6 @@
 import textwrap
 import uuid
 from multiprocessing import queues
-from queue import Empty
 from typing import AnyStr, List, Optional, Union
 
 import pexpect.fdpexpect
@@ -28,7 +27,6 @@ def __init__(self, *args, **kwargs):
         if 'ctx' not in kwargs:
             kwargs['ctx'] = _ctx
 
-        self.lock = _ctx.Lock()
         super().__init__(*args, **kwargs)
 
     def put(self, obj, **kwargs):
@@ -45,17 +43,6 @@ def put(self, obj, **kwargs):
         except:  # noqa # queue might be closed
             pass
 
-    def get_all(self) -> List[bytes]:
-        res = []
-        with self.lock:
-            while True:
-                try:
-                    res.append(self.get_nowait())
-                except Empty:
-                    break
-
-        return res
-
     def write(self, s: AnyStr):
         self.put(s)
 

pytest-embedded/tests/test_base.py

@@ -496,6 +496,28 @@ def test_expect_unity_test_output_multi_dut_with_illegal_chars(dut):
     assert junit_report[1].find('failure') is not None
 
 
+def test_expect_before_match(testdir):
+    testdir.makepyfile(r"""
+        import pexpect
+
+        def test_expect_before_match(dut):
+            dut.write('this would be redirected')
+
+            res = dut.expect('would', return_what_before_match=True)
+            assert res == b'this '
+
+            res = dut.expect_exact('be ', return_what_before_match=True)
+            assert res == b' '
+
+            res = dut.expect('ected', return_what_before_match=True)
+            assert res == b'redir'
+    """)
+
+    result = testdir.runpytest()
+
+    result.assert_outcomes(passed=1)
+
+
 def test_duplicate_stdout_popen(testdir):
     testdir.makepyfile(r"""
         import pytest