modify remove function, pass delete function to each add method

neozwu · neozwu · commit c05b45b1815e · 2017-04-03T15:48:28.000-07:00
diff --git a/google-cloud-compute/src/main/java/com/google/cloud/compute/testing/ResourceCleaner.java b/google-cloud-compute/src/main/java/com/google/cloud/compute/testing/ResourceCleaner.java
@@ -27,7 +27,7 @@
 import com.google.common.base.Function;
 import java.util.HashMap;
 import java.util.Map;
-import sun.reflect.generics.reflectiveObjects.NotImplementedException;
+import java.util.NoSuchElementException;
 
 public class ResourceCleaner {
   private final Compute compute;
@@ -46,13 +46,9 @@ private class ManagedResource<T> {
     final T res;
     Function<T, Void> func;
 
-    private ManagedResource(T res) {
+    private ManagedResource(T res, Function<T, Void> func) {
       this.res = res;
-
-      @SuppressWarnings("unchecked")
-      Function<T, Void> localFunc = (Function<T, Void>) getDeleteMethod(res);
-
-      this.func = localFunc;
+      this.func = func;
     }
 
     public void delete() {
@@ -69,142 +65,115 @@ public void cleanUp() {
   }
 
   public ResourceCleaner add(AddressId addressId) {
-    ManagedResource<AddressId> r = new ManagedResource<>(addressId);
+    ManagedResource<AddressId> r =
+        new ManagedResource<>(
+            addressId,
+            new Function<AddressId, Void>() {
+              @Override
+              public Void apply(AddressId addressId) {
+                compute.deleteAddress(addressId);
+                return null;
+              }
+            });
     resources.put(addressId, r);
     return this;
   }
 
-  public ResourceCleaner remove(AddressId addressId) {
-    resources.remove(addressId);
-    return this;
-  }
-
   public ResourceCleaner add(DiskId diskId) {
-    ManagedResource<DiskId> r = new ManagedResource<>(diskId);
+    ManagedResource<DiskId> r =
+        new ManagedResource<>(
+            diskId,
+            new Function<DiskId, Void>() {
+              @Override
+              public Void apply(DiskId diskId) {
+                compute.deleteDisk(diskId);
+                return null;
+              }
+            });
     resources.put(diskId, r);
     return this;
   }
 
-  public ResourceCleaner remove(DiskId diskId) {
-    resources.remove(diskId);
-    return this;
-  }
-
   public ResourceCleaner add(SnapshotId snapshotId) {
-    ManagedResource<SnapshotId> r = new ManagedResource<>(snapshotId);
+    ManagedResource<SnapshotId> r =
+        new ManagedResource<>(
+            snapshotId,
+            new Function<SnapshotId, Void>() {
+              @Override
+              public Void apply(SnapshotId snapshotId) {
+                compute.deleteSnapshot(snapshotId);
+                return null;
+              }
+            });
     resources.put(snapshotId, r);
     return this;
   }
 
-  public ResourceCleaner remove(SnapshotId snapshotId) {
-    resources.remove(snapshotId);
-    return this;
-  }
-
   public ResourceCleaner add(NetworkId networkId) {
-    ManagedResource<NetworkId> r = new ManagedResource<>(networkId);
+    ManagedResource<NetworkId> r =
+        new ManagedResource<>(
+            networkId,
+            new Function<NetworkId, Void>() {
+              @Override
+              public Void apply(NetworkId networkId) {
+                compute.deleteNetwork(networkId);
+                return null;
+              }
+            });
     resources.put(networkId, r);
     return this;
   }
 
-  public ResourceCleaner remove(NetworkId networkId) {
-    resources.remove(networkId);
-    return this;
-  }
-
   public ResourceCleaner add(SubnetworkId subnetworkId) {
-    ManagedResource<SubnetworkId> r = new ManagedResource<>(subnetworkId);
+    ManagedResource<SubnetworkId> r =
+        new ManagedResource<>(
+            subnetworkId,
+            new Function<SubnetworkId, Void>() {
+              @Override
+              public Void apply(SubnetworkId subnetworkId) {
+                compute.deleteSubnetwork(subnetworkId);
+                return null;
+              }
+            });
     resources.put(subnetworkId, r);
     return this;
   }
 
-  public ResourceCleaner remove(SubnetworkId subnetworkId) {
-    resources.remove(subnetworkId);
-    return this;
-  }
-
   public ResourceCleaner add(ImageId imageId) {
-    ManagedResource<ImageId> r = new ManagedResource<>(imageId);
+    ManagedResource<ImageId> r =
+        new ManagedResource<>(
+            imageId,
+            new Function<ImageId, Void>() {
+              @Override
+              public Void apply(ImageId imageId) {
+                compute.deleteImage(imageId);
+                return null;
+              }
+            });
     resources.put(imageId, r);
     return this;
   }
 
-  public ResourceCleaner remove(ImageId imageId) {
-    resources.remove(imageId);
-    return this;
-  }
-
   public ResourceCleaner add(InstanceId instanceId) {
-    ManagedResource<InstanceId> r = new ManagedResource<>(instanceId);
+    ManagedResource<InstanceId> r =
+        new ManagedResource<>(
+            instanceId,
+            new Function<InstanceId, Void>() {
+              @Override
+              public Void apply(InstanceId instanceId) {
+                compute.deleteInstance(instanceId);
+                return null;
+              }
+            });
     resources.put(instanceId, r);
     return this;
   }
 
-  public ResourceCleaner remove(InstanceId instanceId) {
-    resources.remove(instanceId);
-    return this;
-  }
-
-  private <T> Function<?, Void> getDeleteMethod(T res) {
-
-    if (res instanceof AddressId) {
-      return new Function<AddressId, Void>() {
-        @Override
-        public Void apply(AddressId addressId) {
-          compute.deleteAddress(addressId);
-          return null;
-        }
-      };
-    } else if (res instanceof DiskId) {
-      return new Function<DiskId, Void>() {
-        @Override
-        public Void apply(DiskId diskId) {
-          compute.deleteDisk(diskId);
-          return null;
-        }
-      };
-    } else if (res instanceof SnapshotId) {
-      return new Function<SnapshotId, Void>() {
-        @Override
-        public Void apply(SnapshotId snapshotId) {
-          compute.deleteSnapshot(snapshotId);
-          return null;
-        }
-      };
-    } else if (res instanceof NetworkId) {
-      return new Function<NetworkId, Void>() {
-        @Override
-        public Void apply(NetworkId networkId) {
-          compute.deleteNetwork(networkId);
-          return null;
-        }
-      };
-    } else if (res instanceof SubnetworkId) {
-      return new Function<SubnetworkId, Void>() {
-        @Override
-        public Void apply(SubnetworkId subnetworkId) {
-          compute.deleteSubnetwork(subnetworkId);
-          return null;
-        }
-      };
-    } else if (res instanceof ImageId) {
-      return new Function<ImageId, Void>() {
-        @Override
-        public Void apply(ImageId imageId) {
-          compute.deleteImage(imageId);
-          return null;
-        }
-      };
-    } else if (res instanceof InstanceId) {
-      return new Function<InstanceId, Void>() {
-        @Override
-        public Void apply(InstanceId instanceId) {
-          compute.deleteInstance(instanceId);
-          return null;
-        }
-      };
-    } else {
-      throw new NotImplementedException();
+  public <T> ResourceCleaner remove(T resourceId) {
+    if (!resources.containsKey(resourceId)) {
+        throw new NoSuchElementException(resourceId + " has not been added to managed resources");
     }
+    resources.remove(resourceId);
+    return this;
   }
 }
diff --git a/google-cloud-compute/src/test/java/com/google/cloud/compute/it/ITComputeTest.java b/google-cloud-compute/src/test/java/com/google/cloud/compute/it/ITComputeTest.java
@@ -771,10 +771,6 @@ public void testListRegionAddresses() throws InterruptedException, TimeoutExcept
       count++;
     }
     assertEquals(2, count);
-    compute.deleteAddress(firstAddressId);
-    resourceCleaner.remove(firstAddressId);
-    compute.deleteAddress(secondAddressId);
-    resourceCleaner.remove(secondAddressId);
   }
 
   @Test
@@ -806,10 +802,6 @@ public void testAggregatedListAddresses() throws InterruptedException, TimeoutEx
       count++;
     }
     assertEquals(2, count);
-    compute.deleteAddress(firstAddressId);
-    resourceCleaner.remove(firstAddressId);
-    compute.deleteAddress(secondAddressId);
-    resourceCleaner.remove(secondAddressId);
   }
 
   @Test
@@ -891,10 +883,6 @@ public void testListGlobalAddresses() throws InterruptedException, TimeoutExcept
       count++;
     }
     assertEquals(2, count);
-    compute.deleteAddress(firstAddressId);
-    resourceCleaner.remove(firstAddressId);
-    compute.deleteAddress(secondAddressId);
-    resourceCleaner.remove(secondAddressId);
   }
 
   @Test
@@ -1192,14 +1180,6 @@ public void testListDisksAndSnapshots() throws InterruptedException, TimeoutExce
       count++;
     }
     assertEquals(2, count);
-    compute.deleteDisk(firstDiskId);
-    resourceCleaner.remove(firstDiskId);
-    compute.deleteDisk(secondDiskId);
-    resourceCleaner.remove(secondDiskId);
-    compute.deleteSnapshot(firstSnapshotId);
-    resourceCleaner.remove(firstSnapshotId);
-    compute.deleteSnapshot(secondSnapshotId);
-    resourceCleaner.remove(secondSnapshotId);
   }
 
   @Test
@@ -1239,10 +1219,6 @@ public void testAggregatedListDisks() throws InterruptedException, TimeoutExcept
       count++;
     }
     assertEquals(2, count);
-    compute.deleteDisk(firstDiskId);
-    resourceCleaner.remove(firstDiskId);
-    compute.deleteDisk(secondDiskId);
-    resourceCleaner.remove(secondDiskId);
   }
 
   @Test
@@ -1679,8 +1655,6 @@ public void testCreateGetAndDeleteInstance() throws InterruptedException, Timeou
     operation.waitFor();
     resourceCleaner.remove(instanceId);
     assertNull(compute.getInstance(instanceId));
-    address.delete();
-    resourceCleaner.remove(addressId);
   }
 
   @Test
@@ -1717,8 +1691,6 @@ public void testStartStopAndResetInstance() throws InterruptedException, Timeout
     remoteInstance = compute.getInstance(instanceId,
         Compute.InstanceOption.fields(Compute.InstanceField.STATUS));
     assertEquals(InstanceInfo.Status.RUNNING, remoteInstance.getStatus());
-    remoteInstance.delete();
-    resourceCleaner.remove(instanceId);
   }
 
   @Test
@@ -1765,8 +1737,6 @@ public void testSetInstanceProperties() throws InterruptedException, TimeoutExce
     operation.waitFor();
     remoteInstance = compute.getInstance(instanceId);
     assertEquals(options, remoteInstance.getSchedulingOptions());
-    remoteInstance.delete();
-    resourceCleaner.remove(instanceId);
   }
 
   @Test
@@ -1817,10 +1787,6 @@ public void testAttachAndDetachDisk() throws InterruptedException, TimeoutExcept
     remoteInstance = compute.getInstance(instanceId);
     assertEquals(1, remoteInstance.getAttachedDisks().size());
     assertEquals("dev0", remoteInstance.getAttachedDisks().get(0).getDeviceName());
-    remoteInstance.delete();
-    compute.deleteDisk(diskId);
-    resourceCleaner.remove(instanceId);
-    resourceCleaner.remove(diskId);
   }
 
   @Test
@@ -1865,9 +1831,5 @@ public void testAddAndRemoveAccessConfig() throws InterruptedException, TimeoutE
     instanceOperation.waitFor();
     remoteInstance = compute.getInstance(instanceId);
     assertTrue(remoteInstance.getNetworkInterfaces().get(0).getAccessConfigurations().isEmpty());
-    remoteInstance.delete();
-    remoteAddress.delete();
-    resourceCleaner.remove(instanceId);
-    resourceCleaner.remove(addressId);
   }
 }
