From d245346df56b0e7fa13fc2ad91c2a7c029d82826 Mon Sep 17 00:00:00 2001
From: Ronald Braunstein <rbraunstein@google.com>
Date: Thu, 18 Apr 2024 09:18:29 -0700
Subject: [PATCH] Add TestOnlyProvider for test_module_config

This one is trivial as these modules are always tests and
top-level-test-targets and don't have any complicated struct composition

This requires aosp/3045317 for "m all_teams" to work for
test_module_config_host.

Test: go test ./tradefed_modules
Test: m all_teams
Change-Id: I229e7690088e996452784183a852325c3003ee93
---
 tradefed_modules/test_module_config.go      |  8 ++-
 tradefed_modules/test_module_config_test.go | 63 +++++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/tradefed_modules/test_module_config.go b/tradefed_modules/test_module_config.go
index 686753713..9127f67a9 100644
--- a/tradefed_modules/test_module_config.go
+++ b/tradefed_modules/test_module_config.go
@@ -297,10 +297,16 @@ func (m *testModuleConfigModule) validateBase(ctx android.ModuleContext, depTag
 //  1. manifest file to testcases dir
 //  2. New Module.config / AndroidTest.xml file with our options.
 func (m *testModuleConfigModule) generateManifestAndConfig(ctx android.ModuleContext) {
+	// Keep before early returns.
+	android.SetProvider(ctx, android.TestOnlyProviderKey, android.TestModuleInformation{
+		TestOnly:       true,
+		TopLevelTarget: true,
+	})
+
 	if !m.validateTestSuites(ctx) {
 		return
 	}
-	// Ensure the provider is accurate
+	// Ensure the base provider is accurate
 	if m.provider.TestConfig == nil {
 		return
 	}
diff --git a/tradefed_modules/test_module_config_test.go b/tradefed_modules/test_module_config_test.go
index 41dd3d479..6997228ce 100644
--- a/tradefed_modules/test_module_config_test.go
+++ b/tradefed_modules/test_module_config_test.go
@@ -19,6 +19,8 @@ import (
 	"strconv"
 	"strings"
 	"testing"
+
+	"github.com/google/blueprint"
 )
 
 const bp = `
@@ -347,6 +349,67 @@ func TestModuleConfigHostDuplicateTestSuitesGiveErrors(t *testing.T) {
 		RunTestWithBp(t, badBp)
 }
 
+func TestTestOnlyProvider(t *testing.T) {
+	t.Parallel()
+	ctx := android.GroupFixturePreparers(
+		java.PrepareForTestWithJavaDefaultModules,
+		android.FixtureRegisterWithContext(RegisterTestModuleConfigBuildComponents),
+	).RunTestWithBp(t, `
+                // These should be test-only
+                test_module_config_host {
+                        name: "host-derived-test",
+                        base: "host-base",
+                        exclude_filters: ["android.test.example.devcodelab.DevCodelabTest#testHelloFail"],
+                        include_annotations: ["android.platform.test.annotations.LargeTest"],
+                        test_suites: ["general-tests"],
+                }
+
+                test_module_config {
+                        name: "derived-test",
+                        base: "base",
+                        exclude_filters: ["android.test.example.devcodelab.DevCodelabTest#testHelloFail"],
+                        include_annotations: ["android.platform.test.annotations.LargeTest"],
+                        test_suites: ["general-tests"],
+                }
+
+		android_test {
+			name: "base",
+			sdk_version: "current",
+                        data: ["data/testfile"],
+		}
+
+		java_test_host {
+			name: "host-base",
+                        srcs: ["a.java"],
+                        test_suites: ["general-tests"],
+		}`,
+	)
+
+	// Visit all modules and ensure only the ones that should
+	// marked as test-only are marked as test-only.
+
+	actualTestOnly := []string{}
+	ctx.VisitAllModules(func(m blueprint.Module) {
+		if provider, ok := android.OtherModuleProvider(ctx.TestContext.OtherModuleProviderAdaptor(), m, android.TestOnlyProviderKey); ok {
+			if provider.TestOnly {
+				actualTestOnly = append(actualTestOnly, m.Name())
+			}
+		}
+	})
+	expectedTestOnlyModules := []string{
+		"host-derived-test",
+		"derived-test",
+		// android_test and java_test_host are tests too.
+		"host-base",
+		"base",
+	}
+
+	notEqual, left, right := android.ListSetDifference(expectedTestOnlyModules, actualTestOnly)
+	if notEqual {
+		t.Errorf("test-only: Expected but not found: %v, Found but not expected: %v", left, right)
+	}
+}
+
 // Use for situations where the entries map contains pairs:  [srcPath:installedPath1, srcPath2:installedPath2]
 // and we want to compare the RHS of the pairs, i.e. installedPath1, installedPath2
 func assertEntryPairValues(t *testing.T, actual []string, expected []string) {