1 files changed, 201 insertions, 0 deletions
diff --git a/python/testSrc/com/jetbrains/env/PyEnvTestCase.java b/python/testSrc/com/jetbrains/env/PyEnvTestCase.java
new file mode 100644
index 000000000000..853a5504de63
--- /dev/null
+++ b/python/testSrc/com/jetbrains/env/PyEnvTestCase.java
@@ -0,0 +1,201 @@
+package com.jetbrains.env;
+
+import com.google.common.collect.Lists;
+import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.diagnostic.Logger;
+import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.testFramework.UsefulTestCase;
+import com.intellij.util.SystemProperties;
+import com.intellij.util.ui.UIUtil;
+import com.jetbrains.python.fixtures.PyTestCase;
+import org.hamcrest.Matchers;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.junit.Assume;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.List;
+
+/**
+ * @author traff
+ */
+public abstract class PyEnvTestCase extends UsefulTestCase {
+  private static final Logger LOG = Logger.getInstance(PyEnvTestCase.class.getName());
+
+  private static final String TAGS_FILE = "tags.txt";
+  private static final String PYCHARM_PYTHON_ENVS = "PYCHARM_PYTHON_ENVS";
+  private static final String PYCHARM_PYTHON_VIRTUAL_ENVS = "PYCHARM_PYTHON_VIRTUAL_ENVS";
+
+  protected static final boolean IS_ENV_CONFIGURATION = System.getProperty("pycharm.env") != null;
+
+
+  public static final boolean RUN_REMOTE = SystemProperties.getBooleanProperty("pycharm.run_remote", false);
+
+  public static final boolean RUN_LOCAL = SystemProperties.getBooleanProperty("pycharm.run_local", true);
+
+  /**
+   * Tags that should exist between all tags, available on all interpreters for test to run.
+   * See {@link #PyEnvTestCase(String...)}
+   */
+  @Nullable
+  private final String[] myRequiredTags;
+
+  /**
+   * @param requiredTags tags that should exist on some interpreter for this test to run.
+   *                     if some of these tags do not exist on any interpreter, all test methods would be skipped using
+   *                     {@link org.junit.Assume}.
+   *                     See <a href="http://junit.sourceforge.net/javadoc/org/junit/Assume.html">Assume manual</a>.
+   *                     Check [IDEA-122939] and [TW-25043] as well.
+   */
+  @SuppressWarnings("JUnitTestCaseWithNonTrivialConstructors")
+  protected PyEnvTestCase(@NotNull final String... requiredTags) {
+    myRequiredTags = requiredTags.length > 0 ? requiredTags.clone() : null;
+
+    PyTestCase.initPlatformPrefix();
+  }
+
+  @Override
+  public void setUp() throws Exception {
+    super.setUp();
+    if (myRequiredTags != null) { // Ensure all tags exist between available interpreters
+      Assume.assumeThat(
+        "Can't find some tags between all available interpreter, test (all methods) will be skipped",
+        getAvailableTags(),
+        Matchers.hasItems(myRequiredTags)
+      );
+    }
+  }
+
+  /**
+   * @return all tags available between all interpreters
+   */
+  @NotNull
+  private static Collection<String> getAvailableTags() {
+    final Collection<String> allAvailableTags = new HashSet<String>();
+    for (final String pythonRoot : getPythonRoots()) {
+      allAvailableTags.addAll(loadEnvTags(pythonRoot));
+    }
+    return allAvailableTags;
+  }
+
+  @Override
+  protected void invokeTestRunnable(@NotNull final Runnable runnable) throws Exception {
+    if (runInWriteAction()) {
+      UIUtil.invokeAndWaitIfNeeded(new Runnable() {
+        public void run() {
+          ApplicationManager.getApplication().runWriteAction(runnable);
+        }
+      });
+    }
+    else {
+      runnable.run();
+    }
+  }
+
+  @Override
+  protected boolean runInDispatchThread() {
+    return false;
+  }
+
+  protected boolean runInWriteAction() {
+    return false;
+  }
+
+  public void runPythonTest(final PyTestTask testTask) {
+    runTest(testTask, getTestName(false));
+  }
+
+  public void runTest(@NotNull PyTestTask testTask, @NotNull String testName) {
+    if (notEnvConfiguration()) {
+      fail("Running under teamcity but not by Env configuration. Skipping.");
+      return;
+    }
+
+    List<String> roots = getPythonRoots();
+
+    if (roots.size() == 0) {
+      String msg = testName +
+                   ": environments are not defined. Skipping. \nSpecify either " +
+                   PYCHARM_PYTHON_ENVS +
+                   " or " +
+                   PYCHARM_PYTHON_VIRTUAL_ENVS +
+                   " environment variable.";
+      LOG.warn(msg);
+      System.out.println(msg);
+      return;
+    }
+
+    doRunTests(testTask, testName, roots);
+  }
+
+  protected void doRunTests(PyTestTask testTask, String testName, List<String> roots) {
+    if (RUN_LOCAL) {
+      PyEnvTaskRunner taskRunner = new PyEnvTaskRunner(roots);
+
+      taskRunner.runTask(testTask, testName);
+    }
+  }
+
+
+  public static boolean notEnvConfiguration() {
+    return UsefulTestCase.IS_UNDER_TEAMCITY && !IS_ENV_CONFIGURATION;
+  }
+
+  public static List<String> getPythonRoots() {
+    List<String> roots = Lists.newArrayList();
+
+    String envs = System.getenv(PYCHARM_PYTHON_ENVS);
+    if (envs != null) {
+      roots.addAll(Lists.newArrayList(envs.split(File.pathSeparator)));
+    }
+
+    String virtualEnvs = System.getenv(PYCHARM_PYTHON_VIRTUAL_ENVS);
+
+    if (virtualEnvs != null) {
+      roots.addAll(readVirtualEnvRoots(virtualEnvs));
+    }
+    return roots;
+  }
+
+  protected static List<String> readVirtualEnvRoots(@NotNull String envs) {
+    List<String> result = Lists.newArrayList();
+    String[] roots = envs.split(File.pathSeparator);
+    for (String root : roots) {
+      File virtualEnvRoot = new File(root);
+      File[] virtualenvs = virtualEnvRoot.listFiles();
+      if (virtualenvs != null) {
+        for (File f : virtualenvs) {
+          result.add(f.getAbsolutePath());
+        }
+      } else {
+        LOG.error(root + " is not a directory of doesn't exist");
+      }
+    }
+
+    return result;
+  }
+
+  public static List<String> loadEnvTags(String env) {
+    List<String> envTags;
+
+    try {
+      File parent = new File(env);
+      if (parent.isFile()) {
+        parent = parent.getParentFile();
+      }
+      envTags = com.intellij.openapi.util.io.FileUtil.loadLines(new File(parent, TAGS_FILE));
+    }
+    catch (IOException e) {
+      envTags = Lists.newArrayList();
+    }
+    return envTags;
+  }
+
+  public static String joinStrings(Collection<String> roots, String rootsName) {
+    return roots.size() > 0 ? rootsName + StringUtil.join(roots, ", ") + "\n" : "";
+  }
+}
+