From 612e6ebed710f617d0da006080882c24a5e298dd Mon Sep 17 00:00:00 2001
From: iory <ab.ioryz@gmail.com>
Date: Sat, 18 Apr 2020 15:19:55 +0900
Subject: [PATCH] Add elapsed_time check for test

---
 tests/test_pysnooper.py | 141 ++++++++++++++++++++++++++++++++++------
 tests/utils.py          |  15 ++++-
 2 files changed, 134 insertions(+), 22 deletions(-)

diff --git a/tests/test_pysnooper.py b/tests/test_pysnooper.py
index d4dd5f4..4cdbee5 100644
--- a/tests/test_pysnooper.py
+++ b/tests/test_pysnooper.py
@@ -4,6 +4,7 @@
 import io
 import textwrap
 import threading
+import time
 import types
 import os
 import sys
@@ -51,32 +52,132 @@ def test_string_io():
 
 
 def test_elapsed_time():
-    string_io = io.StringIO()
+    snoop = pysnooper.snoop(elapsed_time=True)
 
-    @pysnooper.snoop(string_io, elapsed_time=True)
-    def my_function(foo):
-        x = 7
-        y = 8
-        return y + x
+    def foo(x):
+        if x == 0:
+            bar1(x)
+            qux()
+            return
 
-    result = my_function('baba')
-    assert result == 15
-    output = string_io.getvalue()
+        with snoop:
+            # There should be line entries for these three lines,
+            # no line entries for anything else in this function,
+            # but calls to all bar functions should be traced
+            foo(x - 1)
+            bar2(x)
+            qux()
+        int(4)
+        bar3(9)
+        return x
+
+    @snoop
+    def bar1(_x):
+        qux()
+
+    @snoop
+    def bar2(_x):
+        qux()
+
+    @snoop
+    def bar3(_x):
+        qux()
+
+    def qux():
+        time.sleep(0.1)
+        return 9  # not traced, mustn't show up
+
+    with mini_toolbox.OutputCapturer(stdout=False,
+                                     stderr=True) as output_capturer:
+        result = foo(2)
+    assert result == 2
+    output = output_capturer.string_io.getvalue()
     assert_output(
         output,
         (
+            # In first with
             SourcePathEntry(),
-            VariableEntry('foo', value_regex="u?'baba'"),
-            CallEntry('def my_function(foo):'),
-            LineEntry('x = 7'),
-            VariableEntry('x', '7'),
-            LineEntry('y = 8'),
-            VariableEntry('y', '8'),
-            LineEntry('return y + x'),
-            ReturnEntry('return y + x'),
-            ReturnValueEntry('15'),
-            ElapsedTimeEntry(),
-        )
+            VariableEntry('x', '2'),
+            VariableEntry('bar1'),
+            VariableEntry('bar2'),
+            VariableEntry('bar3'),
+            VariableEntry('foo'),
+            VariableEntry('qux'),
+            VariableEntry('snoop'),
+            LineEntry('foo(x - 1)'),
+
+            # In with in recursive call
+            VariableEntry('x', '1'),
+            VariableEntry('bar1'),
+            VariableEntry('bar2'),
+            VariableEntry('bar3'),
+            VariableEntry('foo'),
+            VariableEntry('qux'),
+            VariableEntry('snoop'),
+            LineEntry('foo(x - 1)'),
+
+            # Call to bar1 from if block outside with
+            VariableEntry('_x', '0'),
+            VariableEntry('qux'),
+            CallEntry('def bar1(_x):'),
+            LineEntry('qux()'),
+            ReturnEntry('qux()'),
+            ReturnValueEntry('None'),
+            ElapsedTimeEntry(0.1),
+
+            # In with in recursive call
+            LineEntry('bar2(x)'),
+
+            # Call to bar2 from within with
+            VariableEntry('_x', '1'),
+            VariableEntry('qux'),
+            CallEntry('def bar2(_x):'),
+            LineEntry('qux()'),
+            ReturnEntry('qux()'),
+            ReturnValueEntry('None'),
+            ElapsedTimeEntry(0.1),
+
+            # In with in recursive call
+            LineEntry('qux()'),
+            ElapsedTimeEntry(0.4),
+
+            # Call to bar3 from after with
+            VariableEntry('_x', '9'),
+            VariableEntry('qux'),
+            CallEntry('def bar3(_x):'),
+            LineEntry('qux()'),
+            ReturnEntry('qux()'),
+            ReturnValueEntry('None'),
+            ElapsedTimeEntry(0.1),
+
+            # -- Similar to previous few sections,
+            # -- but from first call to foo
+
+            # In with in first call
+            LineEntry('bar2(x)'),
+
+            # Call to bar2 from within with
+            VariableEntry('_x', '2'),
+            VariableEntry('qux'),
+            CallEntry('def bar2(_x):'),
+            LineEntry('qux()'),
+            ReturnEntry('qux()'),
+            ReturnValueEntry('None'),
+            ElapsedTimeEntry(0.1),
+
+            # In with in first call
+            LineEntry('qux()'),
+            ElapsedTimeEntry(0.7),
+
+            # Call to bar3 from after with
+            VariableEntry('_x', '9'),
+            VariableEntry('qux'),
+            CallEntry('def bar3(_x):'),
+            LineEntry('qux()'),
+            ReturnEntry('qux()'),
+            ReturnValueEntry('None'),
+            ElapsedTimeEntry(0.1),
+        ),
     )
 
 
diff --git a/tests/utils.py b/tests/utils.py
index 169d184..2b73555 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -221,8 +221,13 @@ class SourcePathEntry(_BaseValueEntry):
 
 
 class ElapsedTimeEntry(_BasePrintEntry):
-    def __init__(self, prefix=''):
+
+    def __init__(self, elapsed_time_value=None,
+                 tolerance=0.05,
+                 prefix=''):
         _BasePrintEntry.__init__(self, prefix=prefix)
+        self.elapsed_time_value = elapsed_time_value
+        self.tolerance = tolerance
 
     _preamble_pattern = re.compile(
         r"""^Total elapsed time$"""
@@ -232,7 +237,13 @@ class ElapsedTimeEntry(_BasePrintEntry):
         return bool(self._preamble_pattern.match(preamble))
 
     def _check_content(self, content):
-        return True
+        if self.elapsed_time_value:
+            time = pysnooper.pycompat.time_fromisoformat(content)
+            sec = (time.hour * 60 + time.minute) * 60 + time.second + \
+                time.microsecond * (10 ** -6)
+            return abs(sec - self.elapsed_time_value) <= self.tolerance
+        else:
+            return True
 
 
 class _BaseEventEntry(_BaseEntry):