Implement test-only for python_ rules

Similar to: aosp/3040036 and aosp/3035329, implements a test-only
property for python_library and sets it for python_host.

This combined with aosp/3045317 produces:
	gqui from  "flatten(~/aosp-main-with-phones/out/soong/ownership/all_teams.pb, teams)" proto ~/aosp-main-with-phones/build/soong/android/team_proto/team.proto:AllTeams 'select teams.kind, count(*) total, sum(teams.test_only) as test_only, sum(teams.top_level_target) as top_level, where teams.kind like "%py%"  group by teams.kind'
	+-----------------------------+-------+-----------+-----------+
	|         teams.kind          | total | test_only | top_level |
	+-----------------------------+-------+-----------+-----------+
	| cpython3_cc_defaults        |     1 |         0 |         0 |
	| cpython3_cc_prebuilt_binary |     4 |         0 |         0 |
	| python_binary_host          |   222 |         0 |         0 |
	| python_defaults             |    27 |         0 |         0 |
	| python_library              |    38 |         0 |         0 |
	| python_library_host         |   131 |         0 |         0 |
	| python_test                 |     5 |         5 |         5 |
	| python_test_host            |   172 |       172 |       172 |
	+-----------------------------+-------+-----------+-----------+

Change-Id: Iea14fcdbfd782d37ac0a3c780820deeed5f1923d
Test: m blueprint_tests
Test: go test ./python
diff --git a/python/python.go b/python/python.go
index 2b1974e..621e429 100644
--- a/python/python.go
+++ b/python/python.go
@@ -151,6 +151,8 @@
 	// The zip file containing the current module's source/data files, with the
 	// source files precompiled.
 	precompiledSrcsZip android.Path
+
+	sourceProperties android.SourceProperties
 }
 
 // newModule generates new Python base module
@@ -203,7 +205,7 @@
 var _ pythonDependency = (*PythonLibraryModule)(nil)
 
 func (p *PythonLibraryModule) init() android.Module {
-	p.AddProperties(&p.properties, &p.protoProperties)
+	p.AddProperties(&p.properties, &p.protoProperties, &p.sourceProperties)
 	android.InitAndroidArchModule(p, p.hod, p.multilib)
 	android.InitDefaultableModule(p)
 	return p
@@ -421,6 +423,11 @@
 func (p *PythonLibraryModule) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	expandedSrcs := android.PathsForModuleSrcExcludes(ctx, p.properties.Srcs, p.properties.Exclude_srcs)
 	android.SetProvider(ctx, blueprint.SrcsFileProviderKey, blueprint.SrcsFileProviderData{SrcPaths: expandedSrcs.Strings()})
+	// Keep before any early returns.
+	android.SetProvider(ctx, android.TestOnlyProviderKey, android.TestModuleInformation{
+		TestOnly:       Bool(p.sourceProperties.Test_only),
+		TopLevelTarget: p.sourceProperties.Top_level_test_target,
+	})
 
 	// expand data files from "data" property.
 	expandedData := android.PathsForModuleSrc(ctx, p.properties.Data)
diff --git a/python/python_test.go b/python/python_test.go
index 75a6a89..c0b7295 100644
--- a/python/python_test.go
+++ b/python/python_test.go
@@ -18,10 +18,13 @@
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"android/soong/android"
 	"android/soong/cc"
+
+	"github.com/google/blueprint"
 )
 
 type pyModule struct {
@@ -360,6 +363,76 @@
 	}
 }
 
+func TestTestOnlyProvider(t *testing.T) {
+	t.Parallel()
+	ctx := android.GroupFixturePreparers(
+		PrepareForTestWithPythonBuildComponents,
+		android.PrepareForTestWithAllowMissingDependencies,
+	).RunTestWithBp(t, `
+                // These should be test-only
+                python_library { name: "py-lib-test", test_only: true }
+                python_library { name: "py-lib-test-host", test_only: true, host_supported: true }
+                python_test {    name: "py-test", srcs: ["py-test.py"] }
+                python_test_host { name: "py-test-host", srcs: ["py-test-host.py"] }
+                python_binary_host { name: "py-bin-test", srcs: ["py-bin-test.py"] }
+
+                // These should not be.
+                python_library { name: "py-lib" }
+                python_binary_host { name: "py-bin", srcs: ["py-bin.py"] }
+	`)
+
+	// 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{
+		"py-lib-test",
+		"py-lib-test-host",
+		"py-test",
+		"py-test-host",
+	}
+
+	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)
+	}
+}
+
+// Don't allow setting test-only on things that are always tests or never tests.
+func TestInvalidTestOnlyTargets(t *testing.T) {
+	testCases := []string{
+		` python_test { name: "py-test", test_only: true, srcs: ["py-test.py"] } `,
+		` python_test_host { name: "py-test-host", test_only: true, srcs: ["py-test-host.py"] } `,
+		` python_defaults { name: "py-defaults", test_only: true, srcs: ["foo.py"] } `,
+	}
+
+	for i, bp := range testCases {
+		ctx := android.GroupFixturePreparers(
+			PrepareForTestWithPythonBuildComponents,
+			android.FixtureRegisterWithContext(func(ctx android.RegistrationContext) {
+
+				ctx.RegisterModuleType("python_defaults", DefaultsFactory)
+			}),
+			android.PrepareForTestWithAllowMissingDependencies).
+			ExtendWithErrorHandler(android.FixtureIgnoreErrors).
+			RunTestWithBp(t, bp)
+		if len(ctx.Errs) != 1 {
+			t.Errorf("Expected err setting test_only in testcase #%d: %d errs", i, len(ctx.Errs))
+			continue
+		}
+		if !strings.Contains(ctx.Errs[0].Error(), "unrecognized property \"test_only\"") {
+			t.Errorf("ERR: %s bad bp: %s", ctx.Errs[0], bp)
+		}
+	}
+}
+
 func expectModule(t *testing.T, ctx *android.TestContext, name, variant, expectedSrcsZip string, expectedPyRunfiles []string) {
 	module := ctx.ModuleForTests(name, variant)
 
diff --git a/python/test.go b/python/test.go
index 826f353..2b939e7 100644
--- a/python/test.go
+++ b/python/test.go
@@ -36,7 +36,9 @@
 }
 
 func NewTest(hod android.HostOrDeviceSupported) *PythonTestModule {
-	return &PythonTestModule{PythonBinaryModule: *NewBinary(hod)}
+	p := &PythonTestModule{PythonBinaryModule: *NewBinary(hod)}
+	p.sourceProperties = android.SourceProperties{Test_only: proptools.BoolPtr(true), Top_level_test_target: true}
+	return p
 }
 
 func PythonTestHostFactory() android.Module {