public void executeInKeyContext(Runnable runnable, Object key) {
        final Object oldKey = currentKeySupplier.get();
        setCurrentKey(key);
        try {
            runnable.run();
        } finally {
            resetCurrentKey(oldKey);
        }
    }
@Test
    void testExecuteInKeyContext() {
        final int oldKey = 1;
        final int newKey = 2;
        // -1 as unset value
        AtomicInteger setKey = new AtomicInteger(-1);
        DefaultStateManager stateManager =
                new DefaultStateManager(
                        () -> oldKey,
                        k -> setKey.set((Integer) k),
                        new MockStreamingRuntimeContext(false, 1, 0),
                        new MockOperatorStateStore());
        stateManager.executeInKeyContext(() -> assertThat(setKey).hasValue(newKey), newKey);
        assertThat(setKey).hasValue(oldKey);
    }