Relax restriction on suffix of apex_set filename to include .capex

Compressed APEX files generated by the build system have the .capex
extension. This CL allows us using those .capex files with their actual
suffix.

Bug: 197258691
Test: manual + unit test added
Change-Id: I79173ef942326b48b5e659f1873d80f12b212339
diff --git a/apex/androidmk.go b/apex/androidmk.go
index 2f10904..94b8116 100644
--- a/apex/androidmk.go
+++ b/apex/androidmk.go
@@ -382,7 +382,7 @@
 				fmt.Fprintln(w, "LOCAL_MODULE_PATH :=", a.installDir.ToMakePath().String())
 				stemSuffix := apexType.suffix()
 				if a.isCompressed {
-					stemSuffix = ".capex"
+					stemSuffix = imageCapexSuffix
 				}
 				fmt.Fprintln(w, "LOCAL_MODULE_STEM :=", name+stemSuffix)
 				fmt.Fprintln(w, "LOCAL_UNINSTALLABLE_MODULE :=", !a.installable())
diff --git a/apex/apex.go b/apex/apex.go
index fbf6a6f..e3edc68 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -1152,9 +1152,10 @@
 
 const (
 	// File extensions of an APEX for different packaging methods
-	imageApexSuffix = ".apex"
-	zipApexSuffix   = ".zipapex"
-	flattenedSuffix = ".flattened"
+	imageApexSuffix  = ".apex"
+	imageCapexSuffix = ".capex"
+	zipApexSuffix    = ".zipapex"
+	flattenedSuffix  = ".flattened"
 
 	// variant names each of which is for a packaging method
 	imageApexType     = "image"
diff --git a/apex/apex_test.go b/apex/apex_test.go
index 5c7c90b..6027f9b 100644
--- a/apex/apex_test.go
+++ b/apex/apex_test.go
@@ -4640,6 +4640,35 @@
 	}
 }
 
+func TestApexSetFilenameOverride(t *testing.T) {
+	testApex(t, `
+		apex_set {
+ 			name: "com.company.android.myapex",
+			apex_name: "com.android.myapex",
+			set: "company-myapex.apks",
+      filename: "com.company.android.myapex.apex"
+		}
+	`).ModuleForTests("com.company.android.myapex", "android_common_com.android.myapex")
+
+	testApex(t, `
+		apex_set {
+ 			name: "com.company.android.myapex",
+			apex_name: "com.android.myapex",
+			set: "company-myapex.apks",
+      filename: "com.company.android.myapex.capex"
+		}
+	`).ModuleForTests("com.company.android.myapex", "android_common_com.android.myapex")
+
+	testApexError(t, `filename should end in .apex or .capex for apex_set`, `
+		apex_set {
+ 			name: "com.company.android.myapex",
+			apex_name: "com.android.myapex",
+			set: "company-myapex.apks",
+      filename: "some-random-suffix"
+		}
+	`)
+}
+
 func TestPrebuiltOverrides(t *testing.T) {
 	ctx := testApex(t, `
 		prebuilt_apex {
diff --git a/apex/builder.go b/apex/builder.go
index 5baa5c0..3177ee0 100644
--- a/apex/builder.go
+++ b/apex/builder.go
@@ -786,7 +786,7 @@
 
 	if apexType == imageApex && (compressionEnabled || a.testOnlyShouldForceCompression()) {
 		a.isCompressed = true
-		unsignedCompressedOutputFile := android.PathForModuleOut(ctx, a.Name()+".capex.unsigned")
+		unsignedCompressedOutputFile := android.PathForModuleOut(ctx, a.Name()+imageCapexSuffix+".unsigned")
 
 		compressRule := android.NewRuleBuilder(pctx, ctx)
 		compressRule.Command().
@@ -800,7 +800,7 @@
 			FlagWithOutput("--output ", unsignedCompressedOutputFile)
 		compressRule.Build("compressRule", "Generate unsigned compressed APEX file")
 
-		signedCompressedOutputFile := android.PathForModuleOut(ctx, a.Name()+".capex")
+		signedCompressedOutputFile := android.PathForModuleOut(ctx, a.Name()+imageCapexSuffix)
 		if ctx.Config().UseRBE() && ctx.Config().IsEnvTrue("RBE_SIGNAPK") {
 			args["outCommaList"] = signedCompressedOutputFile.String()
 		}
diff --git a/apex/prebuilt.go b/apex/prebuilt.go
index 1bb0fb5..c4794dc 100644
--- a/apex/prebuilt.go
+++ b/apex/prebuilt.go
@@ -924,8 +924,8 @@
 
 func (a *ApexSet) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	a.installFilename = a.InstallFilename()
-	if !strings.HasSuffix(a.installFilename, imageApexSuffix) {
-		ctx.ModuleErrorf("filename should end in %s for apex_set", imageApexSuffix)
+	if !strings.HasSuffix(a.installFilename, imageApexSuffix) && !strings.HasSuffix(a.installFilename, imageCapexSuffix) {
+		ctx.ModuleErrorf("filename should end in %s or %s for apex_set", imageApexSuffix, imageCapexSuffix)
 	}
 
 	inputApex := android.OptionalPathForModuleSrc(ctx, a.prebuiltCommonProperties.Selected_apex).Path()