Preventing PSIAE in tests

python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java

@@ -16,6 +16,7 @@
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.codeInsight.completion.CompletionUtil;
+import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.Ref;
 import com.intellij.psi.PsiElement;
@@ -461,6 +462,9 @@ public class PyCallExpressionHelper {
         final PsiPolyVariantReference reference = ((PyReferenceExpression)callee).getReference(resolveContext);
         final List<PyType> members = new ArrayList<PyType>();
         for (PsiElement target : PyUtil.multiResolveTopPriority(reference)) {
+          if (ApplicationManager.getApplication().isUnitTestMode()) {
+            PyPsiUtils.assertValid(target);
+          }
           if (target != null) {
             final Ref<? extends PyType> typeRef = getCallTargetReturnType(call, target, context);
             if (typeRef != null) {

python/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java

@@ -16,6 +16,7 @@
 package com.jetbrains.python.psi.resolve;
 
 import com.google.common.collect.Lists;
+import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.fileTypes.ExtensionFileNameMatcher;
 import com.intellij.openapi.fileTypes.FileNameMatcher;
@@ -34,10 +35,7 @@ import com.intellij.util.containers.HashSet;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.psi.*;
-import com.jetbrains.python.psi.impl.PyBuiltinCache;
-import com.jetbrains.python.psi.impl.PyReferenceExpressionImpl;
-import com.jetbrains.python.psi.impl.PyResolveResultRater;
-import com.jetbrains.python.psi.impl.ResolveResultList;
+import com.jetbrains.python.psi.impl.*;
 import com.jetbrains.python.psi.types.PyModuleType;
 import com.jetbrains.python.psi.types.PyType;
 import org.jetbrains.annotations.NotNull;
@@ -118,6 +116,9 @@ public class ResolveImportUtil {
 
   @NotNull
   public static List<RatedResolveResult> multiResolveImportElement(PyImportElement importElement, @NotNull final QualifiedName qName) {
+    if (ApplicationManager.getApplication().isUnitTestMode()) {
+      PyPsiUtils.assertValid(importElement);
+    }
     final PyStatement importStatement = importElement.getContainingImportStatement();
     if (importStatement instanceof PyFromImportStatement) {
       return resolveNameInFromImport((PyFromImportStatement)importStatement, qName);

python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java

@@ -170,6 +170,9 @@ public class PyExtractMethodUtil {
 
       // replace statements with call
       callElement = replaceElements(elementsRange, callElement);
+      // Both statements are used in finder, so should be valid at this moment
+      PyPsiUtils.assertValid(statement1);
+      PyPsiUtils.assertValid(statement2);
       callElement = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(callElement);
       if (callElement != null) {
         processDuplicates(callElement, generatedMethod, finder, editor);

python/testSrc/com/jetbrains/env/python/dotNet/PyIronPythonTest.java

@@ -1,19 +1,21 @@
 package com.jetbrains.env.python.dotNet;
 
-import com.intellij.openapi.application.ReadAction;
-import com.intellij.openapi.application.Result;
+import com.intellij.openapi.application.ApplicationManager;
 import com.jetbrains.env.PyEnvTestCase;
 import com.jetbrains.python.psi.PyFile;
+import com.jetbrains.python.sdk.InvalidSdkException;
 import org.jetbrains.annotations.NotNull;
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.io.IOException;
+
 /**
  * IronPython.NET specific tests
  *
  * @author Ilya.Kazakevich
  */
-public class PyIronPythonTest extends PyEnvTestCase {
+public final class PyIronPythonTest extends PyEnvTestCase {
 
   /**
    * IronPython tag
@@ -100,17 +102,19 @@ public class PyIronPythonTest extends PyEnvTestCase {
       "System.Web",
       "import_system.py",
       null
-    );
-    runPythonTest(task);
-    final PyFile skeleton = task.getGeneratedSkeleton();
-    new ReadAction() {
+    ) {
       @Override
-      protected void run(@NotNull Result result) throws Throwable {
-        Assert.assertNotNull("System.Web does not contain class AspNetHostingPermissionLevel. Error generating stub? It has classes  " +
-                             skeleton.getTopLevelClasses(),
-                             skeleton.findTopLevelClass("AspNetHostingPermissionLevel"));
+      public void runTestOn(@NotNull final String sdkHome) throws IOException, InvalidSdkException {
+        super.runTestOn(sdkHome);
+        ApplicationManager.getApplication().runReadAction(() -> {
+          final PyFile skeleton = (PyFile)myFixture.getFile();
+          Assert.assertNotNull("System.Web does not contain class AspNetHostingPermissionLevel. Error generating stub? It has classes  " +
+                               skeleton.getTopLevelClasses(),
+                               skeleton.findTopLevelClass("AspNetHostingPermissionLevel"));
+        });
       }
-    }.execute();
+    };
+    runPythonTest(task);
   }
 
   /**

python/testSrc/com/jetbrains/env/python/dotNet/SkeletonTestTask.java

@@ -4,7 +4,6 @@ import com.google.common.collect.Sets;
 import com.intellij.codeInsight.intention.IntentionAction;
 import com.intellij.codeInspection.LocalQuickFix;
 import com.intellij.codeInspection.ex.QuickFixWrapper;
-import com.intellij.openapi.application.Application;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.application.ModalityState;
 import com.intellij.openapi.application.PathManager;
@@ -22,7 +21,6 @@ import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.inspections.quickfix.GenerateBinaryStubsFix;
 import com.jetbrains.python.inspections.unresolvedReference.PyUnresolvedReferencesInspection;
-import com.jetbrains.python.psi.PyFile;
 import com.jetbrains.python.sdk.InvalidSdkException;
 import com.jetbrains.python.sdk.PythonSdkType;
 import com.jetbrains.python.sdkTools.SdkCreationType;
@@ -64,11 +62,11 @@ class SkeletonTestTask extends PyExecutionFixtureTestTask {
   private final String mySourceFileToRunGenerationOn;
   @NotNull
   private final String myUseQuickFixWithThisModuleOnly;
-  private PyFile myGeneratedSkeleton;
 
   /**
    * @param expectedSkeletonFile          if you want test to compare generated result with some file, provide its name.
-   *                                      Pass null if you do not want to compare result with anything (you may do it yourself with {@link #getGeneratedSkeleton()})
+   *                                      Pass null if you do not want to compare result with anything
+   *                                      (you may do it yourself by overwriting {@link #runTestOn(String)}) but <strong>call super</strong>
    * @param moduleNameToBeGenerated       name of module you think we should generate in dotted notation (like "System.Web" or "com.myModule").
    *                                      System will wait for skeleton file for this module to be generated
    * @param sourceFileToRunGenerationOn   Source file where we should run "generate stubs" on. Be sure to place "caret" on appropriate place!
@@ -142,11 +140,12 @@ class SkeletonTestTask extends PyExecutionFixtureTestTask {
         throw new FileComparisonFailure("asd", skeletonText, actual, skeletonFile.getAbsolutePath());
       }
     }
-    myGeneratedSkeleton = (PyFile)myFixture.configureByFile(skeletonFile.getName());
+    myFixture.configureByFile(skeletonFile.getName());
   }
 
   /**
    * Removes strings that starts with "# by generator", because generator version may change
+   *
    * @param textToClean text to remove strings from
    * @return text after cleanup
    */
@@ -167,12 +166,4 @@ class SkeletonTestTask extends PyExecutionFixtureTestTask {
   public Set<String> getTags() {
     return Collections.unmodifiableSet(IRON_TAGS);
   }
-
-  /**
-   * @return File for generated skeleton. Call it after {@link #runTestOn(String)} only!
-   */
-  @NotNull
-  PyFile getGeneratedSkeleton() {
-    return myGeneratedSkeleton;
-  }
 }