public static <T> MasterTriggerRestoreHook<T> wrapHook(
            MasterTriggerRestoreHook<T> hook, ClassLoader userClassLoader) {

        return new WrappedMasterHook<>(hook, userClassLoader);
    }
@Test
    void wrapHook() throws Exception {
        final String id = "id";

        Thread thread = Thread.currentThread();
        final ClassLoader originalClassLoader = thread.getContextClassLoader();
        final ClassLoader userClassLoader = new URLClassLoader(new URL[0]);

        final CompletableFuture<Void> onceRunnableFuture = new CompletableFuture<>();
        final Runnable onceRunnable =
                () -> {
                    assertThat(Thread.currentThread().getContextClassLoader())
                            .isEqualTo(userClassLoader);
                    assertThat(onceRunnableFuture)
                            .withFailMessage("The runnable shouldn't be called multiple times.")
                            .isNotDone();
                    onceRunnableFuture.complete(null);
                };

        final CompletableFuture<Void> getIdentifierFuture = new CompletableFuture<>();
        final CompletableFuture<Void> closeFuture = new CompletableFuture<>();
        final CompletableFuture<Void> restoreCheckpointFuture = new CompletableFuture<>();
        final CompletableFuture<Void> createCheckpointDataSerializerFuture =
                new CompletableFuture<>();
        MasterTriggerRestoreHook<String> hook =
                new MasterTriggerRestoreHook<String>() {
                    @Override
                    public String getIdentifier() {
                        assertThat(Thread.currentThread().getContextClassLoader())
                                .isEqualTo(userClassLoader);

                        assertThat(getIdentifierFuture)
                                .withFailMessage("The method shouldn't be called multiple times.")
                                .isNotDone();
                        getIdentifierFuture.complete(null);
                        return id;
                    }

                    @Override
                    public void reset() {
                        assertThat(Thread.currentThread().getContextClassLoader())
                                .isEqualTo(userClassLoader);
                    }

                    @Override
                    public void close() {
                        assertThat(Thread.currentThread().getContextClassLoader())
                                .isEqualTo(userClassLoader);
                        assertThat(closeFuture)
                                .withFailMessage("The method shouldn't be called multiple times.")
                                .isNotDone();
                        closeFuture.complete(null);
                    }

                    @Nullable
                    @Override
                    public CompletableFuture<String> triggerCheckpoint(
                            long checkpointId, long timestamp, Executor executor) {
                        assertThat(Thread.currentThread().getContextClassLoader())
                                .isEqualTo(userClassLoader);
                        executor.execute(onceRunnable);
                        return null;
                    }

                    @Override
                    public void restoreCheckpoint(
                            long checkpointId, @Nullable String checkpointData) {
                        assertThat(Thread.currentThread().getContextClassLoader())
                                .isEqualTo(userClassLoader);

                        assertThat(checkpointId).isZero();
                        assertThat(checkpointData).isEmpty();
                        assertThat(restoreCheckpointFuture)
                                .withFailMessage("The method shouldn't be called multiple times.")
                                .isNotDone();
                        restoreCheckpointFuture.complete(null);
                    }

                    @Nullable
                    @Override
                    public SimpleVersionedSerializer<String> createCheckpointDataSerializer() {
                        assertThat(Thread.currentThread().getContextClassLoader())
                                .isEqualTo(userClassLoader);

                        assertThat(createCheckpointDataSerializerFuture)
                                .withFailMessage("The method shouldn't be called multiple times.")
                                .isNotDone();
                        createCheckpointDataSerializerFuture.complete(null);
                        return null;
                    }
                };

        MasterTriggerRestoreHook<String> wrapped = MasterHooks.wrapHook(hook, userClassLoader);

        // verify getIdentifier
        wrapped.getIdentifier();

        assertThat(getIdentifierFuture).isCompleted();
        assertThat(thread.getContextClassLoader()).isEqualTo(originalClassLoader);

        // verify triggerCheckpoint and its wrapped executor
        TestExecutor testExecutor = new TestExecutor();
        wrapped.triggerCheckpoint(0L, 0, testExecutor);
        assertThat(thread.getContextClassLoader()).isEqualTo(originalClassLoader);
        assertThat(testExecutor.command).isNotNull();
        testExecutor.command.run();
        assertThat(onceRunnableFuture).isCompleted();
        assertThat(thread.getContextClassLoader()).isEqualTo(originalClassLoader);

        // verify restoreCheckpoint
        wrapped.restoreCheckpoint(0L, "");
        assertThat(restoreCheckpointFuture).isCompleted();
        assertThat(thread.getContextClassLoader()).isEqualTo(originalClassLoader);

        // verify createCheckpointDataSerializer
        wrapped.createCheckpointDataSerializer();
        assertThat(createCheckpointDataSerializerFuture).isCompleted();
        assertThat(thread.getContextClassLoader()).isEqualTo(originalClassLoader);

        // verify close
        wrapped.close();
        assertThat(closeFuture).isCompleted();
        assertThat(thread.getContextClassLoader()).isEqualTo(originalClassLoader);
    }