content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.kafka.streams.state.internals;
  
  import org.apache.kafka.common.header.internals.RecordHeaders;
  import org.apache.kafka.common.serialization.Serde;
  import org.apache.kafka.common.serialization.Serdes;
  import org.apache.kafka.common.serialization.Serializer;
  import org.apache.kafka.common.serialization.StringSerializer;
  import org.apache.kafka.common.serialization.UUIDSerializer;
  import org.apache.kafka.common.utils.Bytes;
  import org.apache.kafka.streams.KeyValue;
  import org.apache.kafka.streams.processor.StateStoreContext;
  import org.apache.kafka.streams.processor.internals.ProcessorRecordContext;
  import org.apache.kafka.streams.query.Position;
  import org.apache.kafka.streams.state.KeyValueIterator;
  import org.apache.kafka.streams.state.KeyValueStore;
  import org.apache.kafka.streams.state.KeyValueStoreTestDriver;
  import org.apache.kafka.streams.state.StoreBuilder;
  import org.apache.kafka.streams.state.Stores;
  
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  
  import java.util.ArrayList;
  import java.util.List;
  import java.util.NoSuchElementException;
  import java.util.UUID;
  
  import static org.apache.kafka.common.utils.Utils.mkEntry;
  import static org.apache.kafka.common.utils.Utils.mkMap;
  import static org.hamcrest.CoreMatchers.is;
  import static org.hamcrest.CoreMatchers.nullValue;
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.junit.jupiter.api.Assertions.assertArrayEquals;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  public class InMemoryKeyValueStoreTest extends AbstractKeyValueStoreTest {
  
      private KeyValueStore<Bytes, byte[]> byteStore;
      private final Serializer<String> stringSerializer = new StringSerializer();
      private final KeyValueStoreTestDriver<Bytes, byte[]> byteStoreDriver = KeyValueStoreTestDriver.create(Bytes.class, byte[].class);
      private InMemoryKeyValueStore inMemoryKeyValueStore;
  
      @BeforeEach
      public void createStringKeyValueStore() {
          super.before();
          final StateStoreContext byteStoreContext = byteStoreDriver.context();
          final StoreBuilder<KeyValueStore<Bytes, byte[]>> storeBuilder = Stores.keyValueStoreBuilder(
              Stores.inMemoryKeyValueStore("in-memory-byte-store"),
              new Serdes.BytesSerde(),
              new Serdes.ByteArraySerde());
          byteStore = storeBuilder.build();
          byteStore.init(byteStoreContext, byteStore);
          this.inMemoryKeyValueStore = getInMemoryStore();
      }
  
      @AfterEach
      public void after() {
          super.after();
          byteStore.close();
          byteStoreDriver.clear();
      }
  
      @SuppressWarnings("unchecked")
      @Override
      protected <K, V> KeyValueStore<K, V> createKeyValueStore(final StateStoreContext context) {
          final StoreBuilder<KeyValueStore<K, V>> storeBuilder = Stores.keyValueStoreBuilder(
              Stores.inMemoryKeyValueStore("my-store"),
              (Serde<K>) context.keySerde(),
              (Serde<V>) context.valueSerde());
  
          final KeyValueStore<K, V> store = storeBuilder.build();
          store.init(context, store);
          return store;
      }
  
      InMemoryKeyValueStore getInMemoryStore() {
          return new InMemoryKeyValueStore("in-memory-store-test");
      }
  
      @Test
      public void shouldRemoveKeysWithNullValues() {
          store.close();
          // Add any entries that will be restored to any store
          // that uses the driver's context ...
          driver.addEntryToRestoreLog(0, "zero");
          driver.addEntryToRestoreLog(1, "one");
          driver.addEntryToRestoreLog(2, "two");
          driver.addEntryToRestoreLog(3, "three");
          driver.addEntryToRestoreLog(0, null);
  
          store = createKeyValueStore(driver.context());
          context.restore(store.name(), driver.restoredEntries());
  
          assertEquals(3, driver.sizeOf(store));
  
          assertThat(store.get(0), nullValue());
      }
  
  
      @Test
      public void shouldReturnKeysWithGivenPrefix() {
  
          final List<KeyValue<Bytes, byte[]>> entries = new ArrayList<>();
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "k1")),
              stringSerializer.serialize(null, "a")));
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "prefix_3")),
              stringSerializer.serialize(null, "b")));
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "k2")),
              stringSerializer.serialize(null, "c")));
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "prefix_2")),
              stringSerializer.serialize(null, "d")));
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "k3")),
              stringSerializer.serialize(null, "e")));
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "prefix_1")),
              stringSerializer.serialize(null, "f")));
  
          byteStore.putAll(entries);
          byteStore.flush();
  
          final List<String> valuesWithPrefix = new ArrayList<>();
          int numberOfKeysReturned = 0;
  
          try (final KeyValueIterator<Bytes, byte[]> keysWithPrefix = byteStore.prefixScan("prefix", stringSerializer)) {
              while (keysWithPrefix.hasNext()) {
                  final KeyValue<Bytes, byte[]> next = keysWithPrefix.next();
                  valuesWithPrefix.add(new String(next.value));
                  numberOfKeysReturned++;
              }
          }
  
          assertThat(numberOfKeysReturned, is(3));
          assertThat(valuesWithPrefix.get(0), is("f"));
          assertThat(valuesWithPrefix.get(1), is("d"));
          assertThat(valuesWithPrefix.get(2), is("b"));
      }
  
      @Test
      public void shouldReturnKeysWithGivenPrefixExcludingNextKeyLargestKey() {
          final List<KeyValue<Bytes, byte[]>> entries = new ArrayList<>();
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "abc")),
              stringSerializer.serialize(null, "f")));
  
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "abcd")),
              stringSerializer.serialize(null, "f")));
  
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "abce")),
              stringSerializer.serialize(null, "f")));
  
          byteStore.putAll(entries);
          byteStore.flush();
  
          try (final KeyValueIterator<Bytes, byte[]> keysWithPrefixAsabcd = byteStore.prefixScan("abcd", stringSerializer)) {
              int numberOfKeysReturned = 0;
  
              while (keysWithPrefixAsabcd.hasNext()) {
                  keysWithPrefixAsabcd.next().key.get();
                  numberOfKeysReturned++;
              }
  
              assertThat(numberOfKeysReturned, is(1));
          }
      }
  
      @SuppressWarnings("resource")
      @Test
      public void shouldReturnUUIDsWithStringPrefix() {
          final List<KeyValue<Bytes, byte[]>> entries = new ArrayList<>();
          final UUIDSerializer uuidSerializer = new UUIDSerializer();
          final UUID uuid1 = UUID.randomUUID();
          final UUID uuid2 = UUID.randomUUID();
          final String prefix = uuid1.toString().substring(0, 4);
          entries.add(new KeyValue<>(
              new Bytes(uuidSerializer.serialize(null, uuid1)),
              stringSerializer.serialize(null, "a")));
          entries.add(new KeyValue<>(
              new Bytes(uuidSerializer.serialize(null, uuid2)),
              stringSerializer.serialize(null, "b")));
  
          byteStore.putAll(entries);
          byteStore.flush();
  
          final List<String> valuesWithPrefix = new ArrayList<>();
          int numberOfKeysReturned = 0;
  
          try (final KeyValueIterator<Bytes, byte[]> keysWithPrefix = byteStore.prefixScan(prefix, stringSerializer)) {
              while (keysWithPrefix.hasNext()) {
                  final KeyValue<Bytes, byte[]> next = keysWithPrefix.next();
                  valuesWithPrefix.add(new String(next.value));
                  numberOfKeysReturned++;
              }
          }
  
          assertThat(numberOfKeysReturned, is(1));
          assertThat(valuesWithPrefix.get(0), is("a"));
      }
  
      @Test
      public void shouldReturnNoKeys() {
          final List<KeyValue<Bytes, byte[]>> entries = new ArrayList<>();
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "a")),
              stringSerializer.serialize(null, "a")));
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "b")),
              stringSerializer.serialize(null, "c")));
          entries.add(new KeyValue<>(
              new Bytes(stringSerializer.serialize(null, "c")),
              stringSerializer.serialize(null, "e")));
          byteStore.putAll(entries);
          byteStore.flush();
  
          int numberOfKeysReturned = 0;
  
          try (final KeyValueIterator<Bytes, byte[]> keysWithPrefix = byteStore.prefixScan("bb", stringSerializer)) {
              while (keysWithPrefix.hasNext()) {
                  keysWithPrefix.next();
                  numberOfKeysReturned++;
              }
          }
  
          assertThat(numberOfKeysReturned, is(0));
      }
  
      @SuppressWarnings("resource")
      @Test
      public void shouldThrowNullPointerIfPrefixKeySerializerIsNull() {
          assertThrows(NullPointerException.class, () -> byteStore.prefixScan("bb", null));
      }
  
      @Test
      public void shouldMatchPositionAfterPut() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
  
          context.setRecordContext(new ProcessorRecordContext(0, 1, 0, "", new RecordHeaders()));
          inMemoryKeyValueStore.put(bytesKey("key1"), bytesValue("value1"));
          context.setRecordContext(new ProcessorRecordContext(0, 2, 0, "", new RecordHeaders()));
          inMemoryKeyValueStore.put(bytesKey("key2"), bytesValue("value2"));
          context.setRecordContext(new ProcessorRecordContext(0, 3, 0, "", new RecordHeaders()));
          inMemoryKeyValueStore.put(bytesKey("key3"), bytesValue("value3"));
  
          final Position expected = Position.fromMap(mkMap(mkEntry("", mkMap(mkEntry(0, 3L)))));
          final Position actual = inMemoryKeyValueStore.getPosition();
          assertEquals(expected, actual);
      }
  
      @Test
      public void iteratorHasNextOnEmptyStoreShouldReturnFalse() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          try (final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all()) {
              assertFalse(iter.hasNext());
          }
      }
  
      @Test
      public void iteratorHasNextOnDeletedEntryShouldReturnFalse() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key"), bytesValue("value"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          assertTrue(iter.hasNext());
          inMemoryKeyValueStore.delete(bytesKey("key"));
          assertFalse(iter.hasNext());
      }
  
      @Test
      public void iteratorHasNextShouldNotAdvanceIterator() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key"), bytesValue("value"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          assertTrue(iter.hasNext());
          assertTrue(iter.hasNext()); // should still point to the first element
      }
  
      @Test
      public void iteratorHasNextShouldReturnTrueIfElementsRemaining() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key1"), bytesValue("value1"));
          inMemoryKeyValueStore.put(bytesKey("key2"), bytesValue("value2"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          inMemoryKeyValueStore.delete(bytesKey("key1"));
          assertTrue(iter.hasNext());
      }
  
      @Test
      public void iteratorNextShouldReturnNextElement() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key"), bytesValue("value"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          final KeyValue<Bytes, byte[]> next = iter.next();
          assertEquals(bytesKey("key"), next.key);
          assertArrayEquals(bytesValue("value"), next.value);
      }
  
      @Test
      public void iteratorNextAfterHasNextShouldReturnNextElement() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key"), bytesValue("value"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          assertTrue(iter.hasNext());
          final KeyValue<Bytes, byte[]> next = iter.next();
          assertEquals(bytesKey("key"), next.key);
          assertArrayEquals(bytesValue("value"), next.value);
      }
  
      @Test
      public void iteratorNextOnEmptyStoreShouldThrowException() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
          assertThrows(NoSuchElementException.class, iter::next);
      }
  
      @Test
      public void iteratorNextShouldThrowExceptionIfRemainingElementsDeleted() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key1"), bytesValue("value1"));
          inMemoryKeyValueStore.put(bytesKey("key2"), bytesValue("value2"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          final KeyValue<Bytes, byte[]> next = iter.next();
          assertEquals(bytesKey("key1"), next.key);
          assertArrayEquals(bytesValue("value1"), next.value);
  
          inMemoryKeyValueStore.delete(bytesKey("key2"));
          assertThrows(NoSuchElementException.class, iter::next);
      }
  
      @Test
      public void iteratorNextShouldSkipDeletedElements() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key1"), bytesValue("value1"));
          inMemoryKeyValueStore.put(bytesKey("key2"), bytesValue("value2"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          inMemoryKeyValueStore.delete(bytesKey("key1"));
          final KeyValue<Bytes, byte[]> next = iter.next();
          assertEquals(bytesKey("key2"), next.key);
          assertArrayEquals(bytesValue("value2"), next.value);
      }
  
      @Test
      public void iteratorNextShouldIterateOverAllElements() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key1"), bytesValue("value1"));
          inMemoryKeyValueStore.put(bytesKey("key2"), bytesValue("value2"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          final KeyValue<Bytes, byte[]> next1 = iter.next();
          assertEquals(bytesKey("key1"), next1.key);
          assertArrayEquals(bytesValue("value1"), next1.value);
  
          final KeyValue<Bytes, byte[]> next2 = iter.next();
          assertEquals(bytesKey("key2"), next2.key);
          assertArrayEquals(bytesValue("value2"), next2.value);
  
          assertThrows(NoSuchElementException.class, iter::next);
      }
  
      @Test
      public void iteratorPeekNextKeyOnEmptyStoreShouldThrowException() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
          assertThrows(NoSuchElementException.class, iter::peekNextKey);
      }
  
      @Test
      public void iteratorPeekNextKeyOnDeletedEntryShouldThrowException() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key"), bytesValue("value"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          assertEquals(bytesKey("key"), iter.peekNextKey());
          inMemoryKeyValueStore.delete(bytesKey("key"));
          assertThrows(NoSuchElementException.class, iter::peekNextKey);
      }
  
      @Test
      public void iteratorPeekNextKeyShouldNotAdvanceIterator() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key"), bytesValue("value"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          assertEquals(bytesKey("key"), iter.peekNextKey());
          assertEquals(bytesKey("key"), iter.peekNextKey());
      }
  
      @Test
      public void iteratorPeekNextKeyShouldSkipDeletedElements() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          inMemoryKeyValueStore.put(bytesKey("key1"), bytesValue("value1"));
          inMemoryKeyValueStore.put(bytesKey("key2"), bytesValue("value2"));
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          inMemoryKeyValueStore.delete(bytesKey("key1"));
          assertEquals(bytesKey("key2"), iter.peekNextKey());
      }
  
      @Test
      public void iteratorShouldThrowIllegalStateExceptionIfAlreadyClosed() {
          inMemoryKeyValueStore.init(context, inMemoryKeyValueStore);
          final KeyValueIterator<Bytes, byte[]> iter = inMemoryKeyValueStore.all();
  
          iter.close();
          assertThrows(IllegalStateException.class, iter::hasNext);
          assertThrows(IllegalStateException.class, iter::next);
          assertThrows(IllegalStateException.class, iter::peekNextKey);
      }
  
      private byte[] bytesValue(final String value) {
          return value.getBytes();
      }
  
      private Bytes bytesKey(final String key) {
          return Bytes.wrap(key.getBytes());
      }
  
  }
  