test: fix tests failing because of Editions

Author: rahul2393

google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceInfo.java

@@ -24,6 +24,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.FieldMask;
 import com.google.spanner.admin.instance.v1.AutoscalingConfig;
+import com.google.spanner.admin.instance.v1.Instance.Edition;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
@@ -37,6 +38,7 @@ public enum InstanceField implements FieldSelector {
     NODE_COUNT("node_count"),
     PROCESSING_UNITS("processing_units"),
     AUTOSCALING_CONFIG("autoscaling_config"),
+    EDITION("edition"),
     LABELS("labels");
 
     static InstanceField[] defaultFieldsToUpdate(InstanceInfo info) {
@@ -116,6 +118,10 @@ public Builder setAutoscalingConfig(AutoscalingConfig autoscalingConfig) {
       throw new UnsupportedOperationException("Unimplemented");
     }
 
+    public Builder setEdition(Edition edition) {
+      throw new UnsupportedOperationException("Unimplemented");
+    }
+
     public abstract Builder setState(State state);
 
     public abstract Builder addLabel(String key, String value);
@@ -132,6 +138,7 @@ static class BuilderImpl extends Builder {
     private int nodeCount;
     private int processingUnits;
     private AutoscalingConfig autoscalingConfig;
+    private Edition edition;
     private State state;
     private Map<String, String> labels;
     private Timestamp updateTime;
@@ -153,6 +160,7 @@ static class BuilderImpl extends Builder {
       this.labels = new HashMap<>(instance.labels);
       this.updateTime = instance.updateTime;
       this.createTime = instance.createTime;
+      this.edition = instance.edition;
     }
 
     @Override
@@ -197,6 +205,12 @@ public BuilderImpl setAutoscalingConfig(AutoscalingConfig autoscalingConfig) {
       return this;
     }
 
+    @Override
+    public BuilderImpl setEdition(Edition edition) {
+      this.edition = edition;
+      return this;
+    }
+
     @Override
     public BuilderImpl setState(State state) {
       this.state = state;
@@ -227,6 +241,7 @@ public InstanceInfo build() {
   private final int nodeCount;
   private final int processingUnits;
   private final AutoscalingConfig autoscalingConfig;
+  private final Edition edition;
   private final State state;
   private final ImmutableMap<String, String> labels;
   private final Timestamp updateTime;
@@ -239,6 +254,7 @@ public InstanceInfo build() {
     this.nodeCount = builder.nodeCount;
     this.processingUnits = builder.processingUnits;
     this.autoscalingConfig = builder.autoscalingConfig;
+    this.edition = builder.edition;
     this.state = builder.state;
     this.labels = ImmutableMap.copyOf(builder.labels);
     this.updateTime = builder.updateTime;
@@ -283,6 +299,10 @@ public AutoscalingConfig getAutoscalingConfig() {
     return autoscalingConfig;
   }
 
+  public Edition getEdition() {
+    return edition;
+  }
+
   /** Returns the current state of the instance. */
   public State getState() {
     return state;
@@ -306,6 +326,7 @@ public String toString() {
         .add("nodeCount", nodeCount)
         .add("processingUnits", processingUnits)
         .add("autoscaling_config", autoscalingConfig)
+        .add("edition", edition)
         .add("state", state)
         .add("labels", labels)
         .add("createTime", createTime)
@@ -328,6 +349,7 @@ public boolean equals(Object o) {
         && nodeCount == that.nodeCount
         && processingUnits == that.processingUnits
         && Objects.equals(autoscalingConfig, that.autoscalingConfig)
+        && edition == that.edition
         && state == that.state
         && Objects.equals(labels, that.labels)
         && Objects.equals(updateTime, that.updateTime)
@@ -343,6 +365,7 @@ public int hashCode() {
         nodeCount,
         processingUnits,
         autoscalingConfig,
+        edition,
         state,
         labels,
         updateTime,
@@ -365,6 +388,9 @@ com.google.spanner.admin.instance.v1.Instance toProto() {
     if (getAutoscalingConfig() != null) {
       builder.setAutoscalingConfig(getAutoscalingConfig());
     }
+    if (getEdition() != null) {
+      builder.setEdition(getEdition());
+    }
     return builder.build();
   }
 

google-cloud-spanner/src/test/java/com/google/cloud/spanner/InstanceAdminClientImplTest.java

@@ -280,12 +280,13 @@ public void createInstance() throws Exception {
     when(rpc.createInstance(
             "projects/" + PROJECT_ID,
             INSTANCE_ID,
-            getInstanceProto().toBuilder().setProcessingUnits(0).build()))
+            getInstanceProto().toBuilder().setProcessingUnits(0).setEdition(com.google.spanner.admin.instance.v1.Instance.Edition.ENTERPRISE_PLUS).build()))
         .thenReturn(rawOperationFuture);
     OperationFuture<Instance, CreateInstanceMetadata> op =
         client.createInstance(
             InstanceInfo.newBuilder(InstanceId.of(PROJECT_ID, INSTANCE_ID))
                 .setInstanceConfigId(InstanceConfigId.of(PROJECT_ID, CONFIG_ID))
+                .setEdition(com.google.spanner.admin.instance.v1.Instance.Edition.ENTERPRISE_PLUS)
                 .setNodeCount(1)
                 .build());
     assertThat(op.isDone()).isTrue();

google-cloud-spanner/src/test/java/com/google/cloud/spanner/IntegrationTestEnv.java

@@ -160,6 +160,7 @@ private void initializeInstance(InstanceId instanceId) throws Exception {
         InstanceInfo.newBuilder(instanceId)
             .setNodeCount(1)
             .setDisplayName("Test instance")
+            .setEdition(com.google.spanner.admin.instance.v1.Instance.Edition.ENTERPRISE_PLUS)
             .setInstanceConfigId(configId)
             .build();
     OperationFuture<Instance, CreateInstanceMetadata> op =

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITInstanceAdminTest.java

@@ -145,37 +145,27 @@ public void updateInstanceWithAutoscalingConfig() throws Exception {
             .setNodeCount(0)
             .setAutoscalingConfig(autoscalingConfig)
             .build();
-    try {
-      OperationFuture<Instance, UpdateInstanceMetadata> op =
-          instanceClient.updateInstance(toUpdate, InstanceInfo.InstanceField.AUTOSCALING_CONFIG);
-      Instance newInstance = op.get();
-      assertThat(newInstance.getAutoscalingConfig()).isEqualTo(autoscalingConfig);
-
-      Instance newInstanceFromGet =
-          instanceClient.getInstance(env.getTestHelper().getInstanceId().getInstance());
-      assertThat(newInstanceFromGet).isEqualTo(newInstance);
-
-      // Revert back to the instance original state.
-      toUpdate =
-          InstanceInfo.newBuilder(instance.getId())
-              .setAutoscalingConfig(null)
-              .setNodeCount(instance.getNodeCount())
-              .build();
-      instanceClient
-          .updateInstance(
-              toUpdate,
-              InstanceInfo.InstanceField.AUTOSCALING_CONFIG,
-              InstanceInfo.InstanceField.NODE_COUNT)
-          .get();
-    } catch (Exception exception) {
-      // TODO: Remove once the client lib supports creating instances with an Edition.
-      if (!exception
-          .getMessage()
-          .contains("The minimum required Edition for this feature is ENTERPRISE.")) {
-        throw exception;
-      }
-      // ignore this error for now.
-    }
+    OperationFuture<Instance, UpdateInstanceMetadata> op =
+        instanceClient.updateInstance(toUpdate, InstanceInfo.InstanceField.AUTOSCALING_CONFIG);
+    Instance newInstance = op.get();
+    assertThat(newInstance.getAutoscalingConfig()).isEqualTo(autoscalingConfig);
+
+    Instance newInstanceFromGet =
+        instanceClient.getInstance(env.getTestHelper().getInstanceId().getInstance());
+    assertThat(newInstanceFromGet).isEqualTo(newInstance);
+
+    // Revert back to the instance original state.
+    toUpdate =
+        InstanceInfo.newBuilder(instance.getId())
+            .setAutoscalingConfig(null)
+            .setNodeCount(instance.getNodeCount())
+            .build();
+    instanceClient
+        .updateInstance(
+            toUpdate,
+            InstanceInfo.InstanceField.AUTOSCALING_CONFIG,
+            InstanceInfo.InstanceField.NODE_COUNT)
+        .get();
   }
 
   @Test