diff --git a/android/test_suites.go b/android/test_suites.go
index 18744f1..39317ec 100644
--- a/android/test_suites.go
+++ b/android/test_suites.go
@@ -17,6 +17,8 @@
 import (
 	"path/filepath"
 	"strings"
+
+	"github.com/google/blueprint"
 )
 
 func init() {
@@ -37,18 +39,24 @@
 	TestSuites() []string
 }
 
+type TestSuiteInfo struct {
+	TestSuites []string
+}
+
+var TestSuiteInfoProvider = blueprint.NewProvider[TestSuiteInfo]()
+
 func (t *testSuiteFiles) GenerateBuildActions(ctx SingletonContext) {
 	files := make(map[string]map[string]InstallPaths)
 
-	ctx.VisitAllModules(func(m Module) {
-		if tsm, ok := m.(TestSuiteModule); ok {
-			for _, testSuite := range tsm.TestSuites() {
+	ctx.VisitAllModuleProxies(func(m ModuleProxy) {
+		if tsm, ok := OtherModuleProvider(ctx, m, TestSuiteInfoProvider); ok {
+			for _, testSuite := range tsm.TestSuites {
 				if files[testSuite] == nil {
 					files[testSuite] = make(map[string]InstallPaths)
 				}
 				name := ctx.ModuleName(m)
 				files[testSuite][name] = append(files[testSuite][name],
-					OtherModuleProviderOrDefault(ctx, tsm, InstallFilesProvider).InstallFiles...)
+					OtherModuleProviderOrDefault(ctx, m, InstallFilesProvider).InstallFiles...)
 			}
 		}
 	})