From 2bc89de9adf98106c49500d4ccd02b2fec5a350d Mon Sep 17 00:00:00 2001
From: nik <Nikolay.Chashnikov@jetbrains.com>
Date: Fri, 30 Jun 2017 14:23:32 +0200
Subject: [PATCH] injection tests: provide more convenient variant of
 InjectedLanguageManager.registerMultiHostInjector

It allows us to simplify a lot of code in tests and get rid of usages of deprecated method.
---
 .../impl/source/tree/injected/MyTestInjector.java   |  9 +--------
 .../lang/injection/InjectedLanguageManager.java     | 13 +++++++++----
 .../tree/injected/InjectedLanguageManagerImpl.java  |  7 +++++++
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/java/java-impl/src/com/intellij/psi/impl/source/tree/injected/MyTestInjector.java b/java/java-impl/src/com/intellij/psi/impl/source/tree/injected/MyTestInjector.java
index d6a34368c38f..c959cccef6df 100644
--- a/java/java-impl/src/com/intellij/psi/impl/source/tree/injected/MyTestInjector.java
+++ b/java/java-impl/src/com/intellij/psi/impl/source/tree/injected/MyTestInjector.java
@@ -190,14 +190,7 @@ public class MyTestInjector {
         return Arrays.asList(XmlAttributeValue.class);
       }
     };
-    InjectedLanguageManager.getInstance(psiManager.getProject()).registerMultiHostInjector(myMultiHostInjector);
-    Disposer.register(parent, new Disposable() {
-      @Override
-      public void dispose() {
-        boolean b = InjectedLanguageManager.getInstance(psiManager.getProject()).unregisterMultiHostInjector(myMultiHostInjector);
-        assert b;
-      }
-    });
+    InjectedLanguageManager.getInstance(psiManager.getProject()).registerMultiHostInjector(myMultiHostInjector, parent);
 
     final LanguageInjector myInjector = new LanguageInjector() {
       @Override
diff --git a/platform/core-api/src/com/intellij/lang/injection/InjectedLanguageManager.java b/platform/core-api/src/com/intellij/lang/injection/InjectedLanguageManager.java
index 2676a54667b4..8ec9f2cc6fb9 100644
--- a/platform/core-api/src/com/intellij/lang/injection/InjectedLanguageManager.java
+++ b/platform/core-api/src/com/intellij/lang/injection/InjectedLanguageManager.java
@@ -17,6 +17,7 @@
 package com.intellij.lang.injection;
 
 import com.intellij.injected.editor.DocumentWindow;
+import com.intellij.openapi.Disposable;
 import com.intellij.openapi.components.ServiceManager;
 import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.extensions.ExtensionPointName;
@@ -29,6 +30,7 @@ import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiLanguageInjectionHost;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.TestOnly;
 
 import java.util.List;
 
@@ -52,15 +54,18 @@ public abstract class InjectedLanguageManager {
   public abstract int injectedToHost(@NotNull PsiElement injectedContext, int injectedOffset);
 
   /**
-   * Test-only method.
-   * @see com.intellij.lang.injection.MultiHostInjector#MULTIHOST_INJECTOR_EP_NAME
+   * @deprecated use {@link com.intellij.lang.injection.MultiHostInjector#MULTIHOST_INJECTOR_EP_NAME extension point} for production and
+   * {@link #registerMultiHostInjector(MultiHostInjector, Disposable)} for tests
    */
   @Deprecated
   public abstract void registerMultiHostInjector(@NotNull MultiHostInjector injector);
 
+  @TestOnly
+  public abstract void registerMultiHostInjector(@NotNull MultiHostInjector injector, @NotNull Disposable parentDisposable);
+
   /**
-   * Test-only method.
-   * @see com.intellij.lang.injection.MultiHostInjector#MULTIHOST_INJECTOR_EP_NAME
+   * @deprecated use {@link com.intellij.lang.injection.MultiHostInjector#MULTIHOST_INJECTOR_EP_NAME extension point} for production and
+   * {@link #registerMultiHostInjector(MultiHostInjector, Disposable)} for tests
    */
   @Deprecated
   public abstract boolean unregisterMultiHostInjector(@NotNull MultiHostInjector injector);
diff --git a/platform/lang-impl/src/com/intellij/psi/impl/source/tree/injected/InjectedLanguageManagerImpl.java b/platform/lang-impl/src/com/intellij/psi/impl/source/tree/injected/InjectedLanguageManagerImpl.java
index 36d3f3c636cc..be7da928cdb6 100644
--- a/platform/lang-impl/src/com/intellij/psi/impl/source/tree/injected/InjectedLanguageManagerImpl.java
+++ b/platform/lang-impl/src/com/intellij/psi/impl/source/tree/injected/InjectedLanguageManagerImpl.java
@@ -40,6 +40,7 @@ import com.intellij.openapi.progress.ProgressIndicator;
 import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.project.DumbService;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Disposer;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.Segment;
 import com.intellij.openapi.util.TextRange;
@@ -251,6 +252,12 @@ public class InjectedLanguageManagerImpl extends InjectedLanguageManager impleme
     clearInjectorCache();
   }
 
+  @Override
+  public void registerMultiHostInjector(@NotNull MultiHostInjector injector, @NotNull Disposable parentDisposable) {
+    registerMultiHostInjector(injector);
+    Disposer.register(parentDisposable, () -> unregisterMultiHostInjector(injector));
+  }
+
   @Override
   public boolean unregisterMultiHostInjector(@NotNull MultiHostInjector injector) {
     try {