From fa0712573dc80a15fc74bea6aaf52a2d7b46b660 Mon Sep 17 00:00:00 2001
From: Mridula <66699525+mpeddada1@users.noreply.github.com>
Date: Fri, 1 Apr 2022 15:55:21 -0400
Subject: [PATCH] fix(java): register test class for reflection to fix native
 image test (#766)

* fix(java): register test class for reflection to fix native image test
---
 .../features/clients/SpannerFeature.java       | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/java-core/native-image-support/src/main/java/com/google/cloud/nativeimage/features/clients/SpannerFeature.java b/java-core/native-image-support/src/main/java/com/google/cloud/nativeimage/features/clients/SpannerFeature.java
index 7c9fca142163..956f368dacf6 100644
--- a/java-core/native-image-support/src/main/java/com/google/cloud/nativeimage/features/clients/SpannerFeature.java
+++ b/java-core/native-image-support/src/main/java/com/google/cloud/nativeimage/features/clients/SpannerFeature.java
@@ -29,6 +29,7 @@ final class SpannerFeature implements Feature {
 
   private static final String SPANNER_CLASS = "com.google.spanner.v1.SpannerGrpc";
   private static final String SPANNER_TEST_CLASS = "com.google.cloud.spanner.GceTestEnvConfig";
+  private static final String MOCK_CLASS = "com.google.cloud.spanner.MockDatabaseAdminServiceImpl";
   private static final String CLIENT_SIDE_IMPL_CLASS =
       "com.google.cloud.spanner.connection.ClientSideStatementImpl";
   private static final String CLIENT_SIDE_VALUE_CONVERTER =
@@ -50,10 +51,7 @@ final class SpannerFeature implements Feature {
 
   @Override
   public void beforeAnalysis(BeforeAnalysisAccess access) {
-    Class<?> spannerTestClass = access.findClassByName(SPANNER_TEST_CLASS);
-    if (spannerTestClass != null) {
-      NativeImageUtils.registerConstructorsForReflection(access, SPANNER_TEST_CLASS);
-    }
+    registerSpannerTestClasses(access);
     if (access.findClassByName(CLIENT_SIDE_IMPL_CLASS) != null) {
       NativeImageUtils.registerClassHierarchyForReflection(access, CLIENT_SIDE_IMPL_CLASS);
     }
@@ -103,4 +101,16 @@ public void beforeAnalysis(BeforeAnalysisAccess access) {
           "\\Qcom/google/cloud/spanner/connection/ITSqlScriptTest_TestQueryOptions.sql\\E");
     }
   }
+
+  private void registerSpannerTestClasses(BeforeAnalysisAccess access) {
+    Class<?> spannerTestClass = access.findClassByName(SPANNER_TEST_CLASS);
+    if (spannerTestClass != null) {
+      NativeImageUtils.registerConstructorsForReflection(access, SPANNER_TEST_CLASS);
+    }
+    Class<?> mockClass = access.findClassByName(MOCK_CLASS);
+    if (mockClass != null) {
+      NativeImageUtils.registerClassForReflection(
+          access, "com.google.cloud.spanner.MockDatabaseAdminServiceImpl$MockBackup");
+    }
+  }
 }