From 58e816a402cf9d7884abbfbdc04277e0f6a9f8a3 Mon Sep 17 00:00:00 2001
From: Marcin Erdmann <erdi84@gmail.com>
Date: Mon, 21 May 2018 22:02:58 +0300
Subject: [PATCH] [groovy] geb: fix type string in completion for field based
 page content

https://github.com/JetBrains/intellij-community/pull/779
---
 .../jetbrains/plugins/groovy/geb/GebUtil.java |  6 +++--
 .../plugins/groovy/geb/GebTestsTest.groovy    | 18 +++++++++++++
 .../plugins/groovy/util/TestUtils.java        | 25 +++++++++++++++++++
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/geb/GebUtil.java b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/geb/GebUtil.java
index 62bb88de8fc5..59bd3354aba4 100644
--- a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/geb/GebUtil.java
+++ b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/geb/GebUtil.java
@@ -114,8 +114,10 @@ public class GebUtil {
     GrLightField field = new GrLightField(pageOrModuleClass, name, objectType, invokedExpression) {
 
       @Override
-      public PsiType getTypeGroovy() {
-        return block.getReturnType();
+      @NotNull
+      public PsiType getType() {
+        PsiType type = block.getReturnType();
+        return type != null ? type : super.getType();
       }
 
       @Override
diff --git a/plugins/groovy/test/org/jetbrains/plugins/groovy/geb/GebTestsTest.groovy b/plugins/groovy/test/org/jetbrains/plugins/groovy/geb/GebTestsTest.groovy
index a5624b8557f0..2bf06728a5fa 100644
--- a/plugins/groovy/test/org/jetbrains/plugins/groovy/geb/GebTestsTest.groovy
+++ b/plugins/groovy/test/org/jetbrains/plugins/groovy/geb/GebTestsTest.groovy
@@ -95,6 +95,24 @@ class PageWithContent extends geb.Page {
     TestUtils.checkCompletionContains(myFixture, "button", "formField()")
   }
 
+  void testContentElementsCompletionType() {
+    myFixture.configureByText("PageWithContent.groovy", """
+class PageWithContent extends geb.Page {
+  static content = {
+    button { \$('button') }
+    formField { String name -> \$('input', name: name) }
+  }
+  
+  def someMethod() {
+    <caret>
+  }
+}
+""")
+
+    TestUtils.checkCompletionType(myFixture, "button", "geb.navigator.Navigator")
+    TestUtils.checkCompletionType(myFixture, "formField", "geb.navigator.Navigator")
+  }
+
   void testContentMethodReturnType() {
     myFixture.configureByText("PageWithContent.groovy", """
 class PageWithContent extends geb.Page {
diff --git a/plugins/groovy/test/org/jetbrains/plugins/groovy/util/TestUtils.java b/plugins/groovy/test/org/jetbrains/plugins/groovy/util/TestUtils.java
index 3860ae8981fd..0731b4fbdd21 100644
--- a/plugins/groovy/test/org/jetbrains/plugins/groovy/util/TestUtils.java
+++ b/plugins/groovy/test/org/jetbrains/plugins/groovy/util/TestUtils.java
@@ -215,6 +215,31 @@ public abstract class TestUtils {
     }
   }
 
+  public static void checkCompletionType(JavaCodeInsightTestFixture fixture, String lookupString, String expectedTypeCanonicalText) {
+    LookupElement[] lookupElements = fixture.completeBasic();
+    PsiType type = null;
+
+    for (LookupElement lookupElement : lookupElements) {
+      if (lookupElement.getLookupString().equals(lookupString)) {
+        PsiElement element = lookupElement.getPsiElement();
+        if (element instanceof PsiField) {
+          type = ((PsiField)element).getType();
+          break;
+        }
+        if (element instanceof PsiMethod) {
+          type = ((PsiMethod)element).getReturnType();
+          break;
+        }
+      }
+    }
+
+    if (type == null) {
+      Assert.fail("No field or method called '" + lookupString + "' found in completion lookup elements");
+    }
+
+    Assert.assertEquals(expectedTypeCanonicalText, type.getCanonicalText());
+  }
+
   public static void checkResolve(PsiFile file, final String ... expectedUnresolved) {
     final List<String> actualUnresolved = new ArrayList<>();