Add test size attribute to IntelliJ IDE info.

--
MOS_MIGRATED_REVID=121612007

4 files changed, 39 insertions, 0 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/ideinfo/AndroidStudioInfoAspect.java b/src/main/java/com/google/devtools/build/lib/ideinfo/AndroidStudioInfoAspect.java
index 56f917640c..4e4ed765a2 100644
--- a/src/main/java/com/google/devtools/build/lib/ideinfo/AndroidStudioInfoAspect.java
+++ b/src/main/java/com/google/devtools/build/lib/ideinfo/AndroidStudioInfoAspect.java
@@ -55,12 +55,15 @@ import com.google.devtools.build.lib.ideinfo.androidstudio.AndroidStudioIdeInfo.
 import com.google.devtools.build.lib.ideinfo.androidstudio.AndroidStudioIdeInfo.LibraryArtifact;
 import com.google.devtools.build.lib.ideinfo.androidstudio.AndroidStudioIdeInfo.RuleIdeInfo;
 import com.google.devtools.build.lib.ideinfo.androidstudio.AndroidStudioIdeInfo.RuleIdeInfo.Kind;
+import com.google.devtools.build.lib.ideinfo.androidstudio.AndroidStudioIdeInfo.TestInfo;
 import com.google.devtools.build.lib.packages.AspectDefinition;
 import com.google.devtools.build.lib.packages.AspectParameters;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.NativeAspectClass;
+import com.google.devtools.build.lib.packages.NonconfigurableAttributeMapper;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.Rule;
+import com.google.devtools.build.lib.packages.TargetUtils;
 import com.google.devtools.build.lib.rules.android.AndroidIdeInfoProvider;
 import com.google.devtools.build.lib.rules.android.AndroidIdeInfoProvider.SourceDirectory;
 import com.google.devtools.build.lib.rules.android.AndroidSdkProvider;
@@ -363,6 +366,17 @@ public class AndroidStudioInfoAspect extends NativeAspectClass implements Config
           androidIdeInfoProvider, dependenciesResult, ideResolveArtifacts));
     }
 
+    // Test rules
+    if (TargetUtils.isTestRule(base.getTarget())) {
+      TestInfo.Builder builder = TestInfo.newBuilder();
+      String attr = NonconfigurableAttributeMapper.of(base.getTarget().getAssociatedRule())
+          .get("size", Type.STRING);
+      if (attr != null) {
+        builder.setSize(attr);
+      }
+      outputBuilder.setTestInfo(builder);
+    }
+
     AndroidStudioInfoFilesProvider provider = providerBuilder.build();
 
     outputBuilder.addAllDependencies(transform(dependenciesResult.deps, LABEL_TO_STRING));
diff --git a/src/main/protobuf/android_studio_ide_info.proto b/src/main/protobuf/android_studio_ide_info.proto
index b1f2833a84..baedca7cdd 100644
--- a/src/main/protobuf/android_studio_ide_info.proto
+++ b/src/main/protobuf/android_studio_ide_info.proto
@@ -84,6 +84,10 @@ message CToolchainIdeInfo {
   repeated string unfiltered_compiler_option = 9;
 }
 
+message TestInfo {
+  string size = 1;
+}
+
 message RuleIdeInfo {
   enum Kind {
     ANDROID_BINARY = 0;
@@ -125,4 +129,6 @@ message RuleIdeInfo {
   CToolchainIdeInfo c_toolchain_ide_info = 13;
 
   string kind_string = 14;
+
+  TestInfo test_info = 15;
 }
diff --git a/src/test/java/com/google/devtools/build/lib/ideinfo/AndroidStudioInfoAspectTest.java b/src/test/java/com/google/devtools/build/lib/ideinfo/AndroidStudioInfoAspectTest.java
index 52fb2950b3..527a32e6ae 100644
--- a/src/test/java/com/google/devtools/build/lib/ideinfo/AndroidStudioInfoAspectTest.java
+++ b/src/test/java/com/google/devtools/build/lib/ideinfo/AndroidStudioInfoAspectTest.java
@@ -468,6 +468,7 @@ public class AndroidStudioInfoAspectTest extends AndroidStudioInfoAspectTestBase
         "java_test(",
         "   name = 'FooBarTest',",
         "   srcs = ['FooBarTest.java'],",
+        "   size = 'large',",
         "   deps = [':foobar'],",
         ")");
     Map<String, RuleIdeInfo> ruleIdeInfos = buildRuleIdeInfo(
@@ -493,6 +494,8 @@ public class AndroidStudioInfoAspectTest extends AndroidStudioInfoAspectTestBase
     );
     assertThat(testInfo.getJavaRuleIdeInfo().getJdeps().getRelativePath())
         .isEqualTo("java/com/google/example/FooBarTest.jdeps");
+
+    assertThat(testInfo.getTestInfo().getSize()).isEqualTo("large");
   }
 
   @Test
diff --git a/src/test/java/com/google/devtools/build/lib/ideinfo/intellij_info.bzl b/src/test/java/com/google/devtools/build/lib/ideinfo/intellij_info.bzl
index 9cc0e22b63..b708333ab0 100644
--- a/src/test/java/com/google/devtools/build/lib/ideinfo/intellij_info.bzl
+++ b/src/test/java/com/google/devtools/build/lib/ideinfo/intellij_info.bzl
@@ -299,6 +299,18 @@ def android_rule_ide_info(target, ctx):
         ),
         ide_resolve_files)
 
+def test_info(target, ctx):
+  """ Build TestInfo """
+  if not is_test_rule(ctx):
+    return None
+  return struct_omit_none(
+           size = ctx.rule.attr.size,
+         )
+
+def is_test_rule(ctx):
+  kind_string = ctx.rule.kind
+  return kind_string.endswith("_test")
+
 def collect_labels(rule_attrs, attrs):
   """ Collect labels from attribute values.
 
@@ -370,6 +382,9 @@ def _aspect_impl(target, ctx):
   (android_rule_ide_info, android_ide_resolve_files) = android_rule_ide_info(target, ctx)
   ide_resolve_files = ide_resolve_files | android_ide_resolve_files
 
+  # Collect test info
+  test_info = test_info(target, ctx)
+
   # Collect information about exports.
   export_deps = set()
   if hasattr(target, "java"):
@@ -392,6 +407,7 @@ def _aspect_impl(target, ctx):
       java_rule_ide_info = java_rule_ide_info,
       android_rule_ide_info = android_rule_ide_info,
       tags = ctx.rule.attr.tags,
+      test_info = test_info,
   )
 
   # Output the ide information file.