1 files changed, 47 insertions, 0 deletions
diff --git a/python/testSrc/com/jetbrains/python/PyAddImportTest.java b/python/testSrc/com/jetbrains/python/PyAddImportTest.java
index 01f83847aa56..f95187aff819 100644
--- a/python/testSrc/com/jetbrains/python/PyAddImportTest.java
+++ b/python/testSrc/com/jetbrains/python/PyAddImportTest.java
@@ -17,8 +17,13 @@ package com.jetbrains.python;
 
 import com.intellij.openapi.application.Result;
 import com.intellij.openapi.command.WriteCommandAction;
+import com.intellij.psi.PsiPolyVariantReference;
 import com.jetbrains.python.codeInsight.imports.AddImportHelper;
+import com.jetbrains.python.fixtures.PyResolveTestCase;
 import com.jetbrains.python.fixtures.PyTestCase;
+import com.jetbrains.python.psi.PyElement;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * @author yole
@@ -53,4 +58,46 @@ public class PyAddImportTest extends PyTestCase {
     }.execute();
     myFixture.checkResultByFile("addImport/" + getTestName(true) + ".after.py");
   }
+
+  // PY-6020
+  public void testLocalFromImport() {
+    doAddLocalImport("foo", "package.module");
+  }
+
+  // PY-6020
+  public void testLocalImport() {
+    doAddLocalImport("module", null);
+  }
+
+  // PY-13668
+  public void testLocalImportInlineFunctionBody() {
+    testLocalImport();
+  }
+
+  // PY-13668
+  public void testLocalImportInlineBranch() {
+    testLocalImport();
+  }
+
+  /**
+   * Add local import statement
+   * @param name      reference name in corresponding import element
+   * @param qualifier if not {@code null} form {@code from qualifier import name} will be used, otherwise {@code import name}
+   */
+  private void doAddLocalImport(@NotNull final String name, @Nullable final String qualifier) {
+    myFixture.configureByFile("addImport/" + getTestName(true) + ".py");
+    new WriteCommandAction(myFixture.getProject(), myFixture.getFile()) {
+      @Override
+      protected void run(Result result) throws Throwable {
+        final PsiPolyVariantReference reference = PyResolveTestCase.findReferenceByMarker(myFixture.getFile());
+        if (qualifier != null) {
+          AddImportHelper.addLocalFromImportStatement((PyElement)reference.getElement(), qualifier, name);
+        }
+        else {
+          AddImportHelper.addLocalImportStatement((PyElement)reference.getElement(), name);
+        }
+      }
+    }.execute();
+    myFixture.checkResultByFile("addImport/" + getTestName(true) + ".after.py");
+  }
 }