FinalizeBundleHandler { public BeamFnApi.InstructionResponse.Builder finalizeBundle(BeamFnApi.InstructionRequest request) throws Exception { String bundleId = request.getFinalizeBundle().getInstructionId(); Collection<CallbackRegistration> callbacks = bundleFinalizationCallbacks.remove(bundleId); if (callbacks == null) { return BeamFnApi.InstructionResponse.newBuilder() .setFinalizeBundle(FinalizeBundleResponse.getDefaultInstance()); } Collection<Exception> failures = new ArrayList<>(); for (CallbackRegistration callback : callbacks) { try { callback.getCallback().onBundleSuccess(); } catch (Exception e) { failures.add(e); } } if (!failures.isEmpty()) { Exception e = new Exception( String.format("Failed to handle bundle finalization for bundle %s.", bundleId)); for (Exception failure : failures) { e.addSuppressed(failure); } throw e; } return BeamFnApi.InstructionResponse.newBuilder() .setFinalizeBundle(FinalizeBundleResponse.getDefaultInstance()); } FinalizeBundleHandler(ExecutorService executorService); void registerCallbacks(String bundleId, Collection<CallbackRegistration> callbacks); BeamFnApi.InstructionResponse.Builder finalizeBundle(BeamFnApi.InstructionRequest request);  }
@Test public void testFinalizationIgnoresMissingBundleIds() throws Exception { FinalizeBundleHandler handler = new FinalizeBundleHandler(Executors.newCachedThreadPool()); assertEquals(SUCCESSFUL_RESPONSE, handler.finalizeBundle(requestFor("test")).build()); }