From 5b6711c43c1008ba05e93301087eef720d899ccd Mon Sep 17 00:00:00 2001
From: Elizaveta Shashkova <Elizaveta.Shashkova@jetbrains.com>
Date: Mon, 30 Jan 2017 18:14:09 +0300
Subject: [PATCH] Add tests for bytecode modification in frame evaluation

---
 .../test_bytecode_modification.py             | 116 ++++++++++++++++++
 .../env/python/PythonDebuggerTest.java        |  22 +++-
 2 files changed, 135 insertions(+), 3 deletions(-)
 create mode 100644 python/helpers/pydev/tests_pydevd_python/test_bytecode_modification.py

diff --git a/python/helpers/pydev/tests_pydevd_python/test_bytecode_modification.py b/python/helpers/pydev/tests_pydevd_python/test_bytecode_modification.py
new file mode 100644
index 000000000000..9aff9b5d4c6c
--- /dev/null
+++ b/python/helpers/pydev/tests_pydevd_python/test_bytecode_modification.py
@@ -0,0 +1,116 @@
+import sys
+import unittest
+from io import StringIO
+
+from _pydevd_frame_eval.pydevd_modify_bytecode import insert_code
+
+TRACE_MESSAGE = "Trace called"
+
+def tracing():
+    print(TRACE_MESSAGE)
+
+
+def bar(a, b):
+    return a + b
+
+
+class TestInsertCode(unittest.TestCase):
+    lines_separator = "---Line tested---"
+
+    def check_insert_every_line(self, func_to_modify, func_to_insert, number_of_lines):
+        first_line = func_to_modify.__code__.co_firstlineno + 1
+        last_line = first_line + number_of_lines
+        for i in range(first_line, last_line):
+            self.check_insert_to_line(func_to_modify, func_to_insert, i)
+            print(self.lines_separator)
+
+    def check_insert_to_line(self, func_to_modify, func_to_insert, line_number):
+        code_orig = func_to_modify.__code__
+        code_to_insert = func_to_insert.__code__
+        result = insert_code(code_orig, code_to_insert, line_number)
+        exec(result)
+        output = sys.stdout.getvalue().strip().split(self.lines_separator)[-1]
+        self.assertTrue(TRACE_MESSAGE in output)
+
+    def test_assignment(self):
+        self.original_stdout = sys.stdout
+        sys.stdout = StringIO()
+
+        try:
+            def original():
+                a = 1
+                b = 2
+                c = 3
+
+            self.check_insert_every_line(original, tracing, 3)
+
+        finally:
+            sys.stdout = self.original_stdout
+
+    def test_for_loop(self):
+        self.original_stdout = sys.stdout
+        sys.stdout = StringIO()
+
+        try:
+            def original():
+                n = 3
+                sum = 0
+                for i in range(n):
+                    sum += i
+                return sum
+
+            self.check_insert_every_line(original, tracing, 5)
+
+        finally:
+            sys.stdout = self.original_stdout
+
+    def test_if_else(self):
+        self.original_stdout = sys.stdout
+        sys.stdout = StringIO()
+
+        try:
+            def original():
+                if True:
+                    a = 1
+                else:
+                    a = 0
+                print(a)
+
+            self.check_insert_to_line(original, tracing, original.__code__.co_firstlineno + 2)
+            self.check_insert_to_line(original, tracing, original.__code__.co_firstlineno + 5)
+
+        finally:
+            sys.stdout = self.original_stdout
+
+    def test_call_other_function(self):
+        self.original_stdout = sys.stdout
+        sys.stdout = StringIO()
+
+        try:
+            def original():
+                a = 1
+                b = 3
+                c = bar(a, b)
+                return c
+
+            self.check_insert_every_line(original, tracing, 4)
+
+        finally:
+            sys.stdout = self.original_stdout
+
+    def test_class_method(self):
+        self.original_stdout = sys.stdout
+        sys.stdout = StringIO()
+
+        try:
+            class A(object):
+                @staticmethod
+                def foo():
+                    print("i'm in foo")
+
+            original = A.foo
+            self.check_insert_to_line(original, tracing, original.__code__.co_firstlineno + 2)
+
+        finally:
+            sys.stdout = self.original_stdout
+
diff --git a/python/testSrc/com/jetbrains/env/python/PythonDebuggerTest.java b/python/testSrc/com/jetbrains/env/python/PythonDebuggerTest.java
index a5126c923cfe..3f7a004aedb2 100644
--- a/python/testSrc/com/jetbrains/env/python/PythonDebuggerTest.java
+++ b/python/testSrc/com/jetbrains/env/python/PythonDebuggerTest.java
@@ -27,6 +27,7 @@ import com.jetbrains.python.debugger.settings.PySteppingFilter;
 import com.jetbrains.python.sdk.flavors.PythonSdkFlavor;
 import com.jetbrains.python.sdkTools.SdkCreationType;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.junit.Test;
 
 import java.util.ArrayList;
@@ -79,16 +80,22 @@ public class PythonDebuggerTest extends PyEnvTestCase {
   @Test
   @Staging
   public void testPydevTests_Debugger() {
-    unittests("tests_pydevd_python/test_debugger.py");
+    unittests("tests_pydevd_python/test_debugger.py", null);
   }
 
   @Test
   @Staging
   public void testPydevMonkey() {
-    unittests("tests_pydevd_python/test_pydev_monkey.py");
+    unittests("tests_pydevd_python/test_pydev_monkey.py", null);
   }
 
-  private void unittests(final String script) {
+  @Test
+  @Staging
+  public void testBytecodeModification() {
+    unittests("tests_pydevd_python/test_bytecode_modification.py", Sets.newHashSet("python36"));
+  }
+
+  private void unittests(final String script, @Nullable Set<String> tags) {
     runPythonTest(new PyProcessWithConsoleTestTask<PyUnitTestProcessRunner>("/helpers/pydev", SdkCreationType.SDK_PACKAGES_ONLY) {
 
       @NotNull
@@ -110,6 +117,15 @@ public class PythonDebuggerTest extends PyEnvTestCase {
                                       @NotNull final String all) {
         runner.assertAllTestsPassed();
       }
+
+      @NotNull
+      @Override
+      public Set<String> getTags() {
+        if (tags == null) {
+          return super.getTags();
+        }
+        return tags;
+      }
     });
   }