Add apex_test for mixed builds

Test: go test soong tests
Change-Id: I33327e1fe7d6a4bbf32890d9fbd5453c28a9c5b1
diff --git a/android/bazel_handler.go b/android/bazel_handler.go
index 1f87410..2524370 100644
--- a/android/bazel_handler.go
+++ b/android/bazel_handler.go
@@ -250,7 +250,8 @@
 }
 
 func (m MockBazelContext) GetApexInfo(label string, _ configKey) (cquery.ApexInfo, error) {
-	panic("unimplemented")
+	result, _ := m.LabelToApexInfo[label]
+	return result, nil
 }
 
 func (m MockBazelContext) GetCcUnstrippedInfo(label string, _ configKey) (cquery.CcUnstrippedInfo, error) {
diff --git a/apex/apex.go b/apex/apex.go
index a3872d3..cb5a5e9 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -1893,10 +1893,13 @@
 	a.outputFile = a.outputApexFile
 	a.setCompression(ctx)
 
+	// TODO(b/257829940): These are used by the apex_keys_text singleton; would probably be a clearer
+	// interface if these were set in a provider rather than the module itself
 	a.publicKeyFile = android.PathForBazelOut(ctx, outputs.BundleKeyInfo[0])
 	a.privateKeyFile = android.PathForBazelOut(ctx, outputs.BundleKeyInfo[1])
 	a.containerCertificateFile = android.PathForBazelOut(ctx, outputs.ContainerKeyInfo[0])
 	a.containerPrivateKeyFile = android.PathForBazelOut(ctx, outputs.ContainerKeyInfo[1])
+
 	apexType := a.properties.ApexType
 	switch apexType {
 	case imageApex:
diff --git a/apex/apex_test.go b/apex/apex_test.go
index 10adf8d..a02300e 100644
--- a/apex/apex_test.go
+++ b/apex/apex_test.go
@@ -29,6 +29,7 @@
 	"github.com/google/blueprint/proptools"
 
 	"android/soong/android"
+	"android/soong/bazel/cquery"
 	"android/soong/bpf"
 	"android/soong/cc"
 	"android/soong/dexpreopt"
@@ -9744,3 +9745,67 @@
 	libcCoreVariant := result.ModuleForTests("libc.apiimport", "android_arm64_armv8-a_shared").Module()
 	android.AssertBoolEquals(t, "core variant should link against source libc", true, hasDep(libfooCoreVariant, libcCoreVariant))
 }
+
+func TestApexImageInMixedBuilds(t *testing.T) {
+	bp := `
+apex_key{
+	name: "foo_key",
+}
+apex {
+	name: "foo",
+	key: "foo_key",
+	updatable: true,
+	min_sdk_version: "31",
+	file_contexts: ":myapex-file_contexts",
+	bazel_module: { label: "//:foo" },
+}`
+
+	outputBaseDir := "out/bazel"
+	result := android.GroupFixturePreparers(
+		prepareForApexTest,
+		android.FixtureModifyConfig(func(config android.Config) {
+			config.BazelContext = android.MockBazelContext{
+				OutputBaseDir: outputBaseDir,
+				LabelToApexInfo: map[string]cquery.ApexInfo{
+					"//:foo": cquery.ApexInfo{
+						SignedOutput:     "signed_out.apex",
+						UnsignedOutput:   "unsigned_out.apex",
+						BundleKeyInfo:    []string{"public_key", "private_key"},
+						ContainerKeyInfo: []string{"container_cert", "container_private"},
+
+						// unused
+						PackageName:  "pkg_name",
+						ProvidesLibs: []string{"a", "b"},
+						RequiresLibs: []string{"c", "d"},
+					},
+				},
+			}
+		}),
+	).RunTestWithBp(t, bp)
+
+	m := result.ModuleForTests("foo", "android_common_foo_image").Module()
+	ab, ok := m.(*apexBundle)
+	if !ok {
+		t.Fatalf("Expected module to be an apexBundle, was not")
+	}
+
+	if w, g := "out/bazel/execroot/__main__/public_key", ab.publicKeyFile.String(); w != g {
+		t.Errorf("Expected public key %q, got %q", w, g)
+	}
+
+	if w, g := "out/bazel/execroot/__main__/private_key", ab.privateKeyFile.String(); w != g {
+		t.Errorf("Expected private key %q, got %q", w, g)
+	}
+
+	if w, g := "out/bazel/execroot/__main__/container_cert", ab.containerCertificateFile.String(); w != g {
+		t.Errorf("Expected public container key %q, got %q", w, g)
+	}
+
+	if w, g := "out/bazel/execroot/__main__/container_private", ab.containerPrivateKeyFile.String(); w != g {
+		t.Errorf("Expected private container key %q, got %q", w, g)
+	}
+
+	if w, g := "out/bazel/execroot/__main__/signed_out.apex", ab.outputFile.String(); w != g {
+		t.Errorf("Expected output file %q, got %q", w, g)
+	}
+}