diff --git a/android/arch_test.go b/android/arch_test.go
index 8525b03..4cef4c8 100644
--- a/android/arch_test.go
+++ b/android/arch_test.go
@@ -353,9 +353,9 @@
 		t.Run(tt.name, func(t *testing.T) {
 			config := TestArchConfig(buildDir, nil, bp, nil)
 
-			ctx := NewTestArchContext()
+			ctx := NewTestArchContext(config)
 			ctx.RegisterModuleType("module", archTestModuleFactory)
-			ctx.Register(config)
+			ctx.Register()
 			if tt.config != nil {
 				tt.config(config)
 			}
@@ -442,9 +442,9 @@
 		t.Run(tt.name, func(t *testing.T) {
 			config := TestArchConfigNativeBridge(buildDir, nil, bp, nil)
 
-			ctx := NewTestArchContext()
+			ctx := NewTestArchContext(config)
 			ctx.RegisterModuleType("module", archTestModuleFactory)
-			ctx.Register(config)
+			ctx.Register()
 			if tt.config != nil {
 				tt.config(config)
 			}