code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.source.relationaldb.state;
  
  import static io.airbyte.integrations.source.relationaldb.state.StateTestConstants.CURSOR;
  import static io.airbyte.integrations.source.relationaldb.state.StateTestConstants.CURSOR_FIELD1;
  import static io.airbyte.integrations.source.relationaldb.state.StateTestConstants.CURSOR_FIELD2;
  import static io.airbyte.integrations.source.relationaldb.state.StateTestConstants.NAMESPACE;
  import static io.airbyte.integrations.source.relationaldb.state.StateTestConstants.NAME_NAMESPACE_PAIR1;
  import static io.airbyte.integrations.source.relationaldb.state.StateTestConstants.NAME_NAMESPACE_PAIR2;
  import static io.airbyte.integrations.source.relationaldb.state.StateTestConstants.STREAM_NAME1;
  import static io.airbyte.integrations.source.relationaldb.state.StateTestConstants.STREAM_NAME2;
  import static io.airbyte.integrations.source.relationaldb.state.StateTestConstants.STREAM_NAME3;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.mockito.Mockito.mock;
  
  import io.airbyte.commons.json.Jsons;
  import io.airbyte.integrations.base.AirbyteStreamNameNamespacePair;
  import io.airbyte.integrations.source.relationaldb.models.DbState;
  import io.airbyte.integrations.source.relationaldb.models.DbStreamState;
  import io.airbyte.protocol.models.AirbyteStateMessage;
  import io.airbyte.protocol.models.AirbyteStateMessage.AirbyteStateType;
  import io.airbyte.protocol.models.AirbyteStream;
  import io.airbyte.protocol.models.AirbyteStreamState;
  import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
  import io.airbyte.protocol.models.ConfiguredAirbyteStream;
  import io.airbyte.protocol.models.StreamDescriptor;
  import java.util.ArrayList;
  import java.util.Comparator;
  import java.util.List;
  import java.util.Optional;
  import java.util.stream.Collectors;
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.api.Test;
  
  /**
   * Test suite for the {@link StreamStateManager} class.
   */
  public class StreamStateManagerTest {
  
    @Test
    void testCreationFromInvalidState() {
      final AirbyteStateMessage airbyteStateMessage = new AirbyteStateMessage()
          .withType(AirbyteStateType.STREAM)
          .withStream(new AirbyteStreamState()
              .withStreamDescriptor(new StreamDescriptor().withName(STREAM_NAME1).withNamespace(NAMESPACE))
              .withStreamState(Jsons.jsonNode("Not a state object")));
      final ConfiguredAirbyteCatalog catalog = mock(ConfiguredAirbyteCatalog.class);
  
      Assertions.assertDoesNotThrow(() -> {
        final StateManager stateManager = new StreamStateManager(List.of(airbyteStateMessage), catalog);
        assertNotNull(stateManager);
      });
    }
  
    @Test
    void testGetters() {
      final List<AirbyteStateMessage> state = new ArrayList<>();
      state.add(createStreamState(STREAM_NAME1, NAMESPACE, List.of(CURSOR_FIELD1), CURSOR));
      state.add(createStreamState(STREAM_NAME2, NAMESPACE, List.of(), null));
  
      final ConfiguredAirbyteCatalog catalog = new ConfiguredAirbyteCatalog()
          .withStreams(List.of(
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME1).withNamespace(NAMESPACE))
                  .withCursorField(List.of(CURSOR_FIELD1)),
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME2).withNamespace(NAMESPACE))));
  
      final StateManager stateManager = new StreamStateManager(state, catalog);
  
      assertEquals(Optional.of(CURSOR_FIELD1), stateManager.getOriginalCursorField(NAME_NAMESPACE_PAIR1));
      assertEquals(Optional.of(CURSOR), stateManager.getOriginalCursor(NAME_NAMESPACE_PAIR1));
      assertEquals(Optional.of(CURSOR_FIELD1), stateManager.getCursorField(NAME_NAMESPACE_PAIR1));
      assertEquals(Optional.of(CURSOR), stateManager.getCursor(NAME_NAMESPACE_PAIR1));
  
      assertEquals(Optional.empty(), stateManager.getOriginalCursorField(NAME_NAMESPACE_PAIR2));
      assertEquals(Optional.empty(), stateManager.getOriginalCursor(NAME_NAMESPACE_PAIR2));
      assertEquals(Optional.empty(), stateManager.getCursorField(NAME_NAMESPACE_PAIR2));
      assertEquals(Optional.empty(), stateManager.getCursor(NAME_NAMESPACE_PAIR2));
    }
  
    @Test
    void testToState() {
      final ConfiguredAirbyteCatalog catalog = new ConfiguredAirbyteCatalog()
          .withStreams(List.of(
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME1).withNamespace(NAMESPACE))
                  .withCursorField(List.of(CURSOR_FIELD1)),
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME2).withNamespace(NAMESPACE))
                  .withCursorField(List.of(CURSOR_FIELD2)),
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME3).withNamespace(NAMESPACE))));
  
      final StateManager stateManager = new StreamStateManager(createDefaultState(), catalog);
  
      final DbState expectedFirstDbState = new DbState()
          .withCdc(false)
          .withStreams(List.of(
              new DbStreamState()
                  .withStreamName(STREAM_NAME1)
                  .withStreamNamespace(NAMESPACE)
                  .withCursorField(List.of(CURSOR_FIELD1))
                  .withCursor("a"),
              new DbStreamState()
                  .withStreamName(STREAM_NAME2)
                  .withStreamNamespace(NAMESPACE)
                  .withCursorField(List.of(CURSOR_FIELD2)),
              new DbStreamState()
                  .withStreamName(STREAM_NAME3)
                  .withStreamNamespace(NAMESPACE))
              .stream().sorted(Comparator.comparing(DbStreamState::getStreamName)).collect(Collectors.toList()));
      final AirbyteStateMessage expectedFirstEmission =
          createStreamState(STREAM_NAME1, NAMESPACE, List.of(CURSOR_FIELD1), "a").withData(Jsons.jsonNode(expectedFirstDbState));
  
      final AirbyteStateMessage actualFirstEmission = stateManager.updateAndEmit(NAME_NAMESPACE_PAIR1, "a");
      assertEquals(expectedFirstEmission, actualFirstEmission);
  
      final DbState expectedSecondDbState = new DbState()
          .withCdc(false)
          .withStreams(List.of(
              new DbStreamState()
                  .withStreamName(STREAM_NAME1)
                  .withStreamNamespace(NAMESPACE)
                  .withCursorField(List.of(CURSOR_FIELD1))
                  .withCursor("a"),
              new DbStreamState()
                  .withStreamName(STREAM_NAME2)
                  .withStreamNamespace(NAMESPACE)
                  .withCursorField(List.of(CURSOR_FIELD2))
                  .withCursor("b"),
              new DbStreamState()
                  .withStreamName(STREAM_NAME3)
                  .withStreamNamespace(NAMESPACE))
              .stream().sorted(Comparator.comparing(DbStreamState::getStreamName)).collect(Collectors.toList()));
      final AirbyteStateMessage expectedSecondEmission =
          createStreamState(STREAM_NAME2, NAMESPACE, List.of(CURSOR_FIELD2), "b").withData(Jsons.jsonNode(expectedSecondDbState));
  
      final AirbyteStateMessage actualSecondEmission = stateManager.updateAndEmit(NAME_NAMESPACE_PAIR2, "b");
      assertEquals(expectedSecondEmission, actualSecondEmission);
    }
  
    @Test
    void testToStateWithoutCursorInfo() {
      final ConfiguredAirbyteCatalog catalog = new ConfiguredAirbyteCatalog()
          .withStreams(List.of(
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME1).withNamespace(NAMESPACE))
                  .withCursorField(List.of(CURSOR_FIELD1)),
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME2).withNamespace(NAMESPACE))
                  .withCursorField(List.of(CURSOR_FIELD2)),
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME3).withNamespace(NAMESPACE))));
      final AirbyteStreamNameNamespacePair airbyteStreamNameNamespacePair = new AirbyteStreamNameNamespacePair("other", "other");
  
      final StateManager stateManager = new StreamStateManager(createDefaultState(), catalog);
      final AirbyteStateMessage airbyteStateMessage = stateManager.toState(Optional.of(airbyteStreamNameNamespacePair));
      assertNotNull(airbyteStateMessage);
      assertEquals(AirbyteStateType.STREAM, airbyteStateMessage.getType());
      assertNotNull(airbyteStateMessage.getStream());
    }
  
    @Test
    void testToStateWithoutStreamPair() {
      final ConfiguredAirbyteCatalog catalog = new ConfiguredAirbyteCatalog()
          .withStreams(List.of(
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME1).withNamespace(NAMESPACE))
                  .withCursorField(List.of(CURSOR_FIELD1)),
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME2).withNamespace(NAMESPACE))
                  .withCursorField(List.of(CURSOR_FIELD2)),
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME3).withNamespace(NAMESPACE))));
  
      final StateManager stateManager = new StreamStateManager(createDefaultState(), catalog);
      final AirbyteStateMessage airbyteStateMessage = stateManager.toState(Optional.empty());
      assertNotNull(airbyteStateMessage);
      assertEquals(AirbyteStateType.STREAM, airbyteStateMessage.getType());
      assertNotNull(airbyteStateMessage.getStream());
      assertNull(airbyteStateMessage.getStream().getStreamState());
    }
  
    @Test
    void testToStateNullCursorField() {
      final ConfiguredAirbyteCatalog catalog = new ConfiguredAirbyteCatalog()
          .withStreams(List.of(
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME1).withNamespace(NAMESPACE))
                  .withCursorField(List.of(CURSOR_FIELD1)),
              new ConfiguredAirbyteStream()
                  .withStream(new AirbyteStream().withName(STREAM_NAME2).withNamespace(NAMESPACE))));
      final StateManager stateManager = new StreamStateManager(createDefaultState(), catalog);
  
      final DbState expectedFirstDbState = new DbState()
          .withCdc(false)
          .withStreams(List.of(
              new DbStreamState()
                  .withStreamName(STREAM_NAME1)
                  .withStreamNamespace(NAMESPACE)
                  .withCursorField(List.of(CURSOR_FIELD1))
                  .withCursor("a"),
              new DbStreamState()
                  .withStreamName(STREAM_NAME2)
                  .withStreamNamespace(NAMESPACE))
              .stream().sorted(Comparator.comparing(DbStreamState::getStreamName)).collect(Collectors.toList()));
  
      final AirbyteStateMessage expectedFirstEmission =
          createStreamState(STREAM_NAME1, NAMESPACE, List.of(CURSOR_FIELD1), "a").withData(Jsons.jsonNode(expectedFirstDbState));
      final AirbyteStateMessage actualFirstEmission = stateManager.updateAndEmit(NAME_NAMESPACE_PAIR1, "a");
      assertEquals(expectedFirstEmission, actualFirstEmission);
    }
  
    @Test
    void testCdcStateManager() {
      final ConfiguredAirbyteCatalog catalog = mock(ConfiguredAirbyteCatalog.class);
      final StateManager stateManager = new StreamStateManager(
          List.of(new AirbyteStateMessage().withType(AirbyteStateType.STREAM).withStream(new AirbyteStreamState())), catalog);
      Assertions.assertThrows(UnsupportedOperationException.class, () -> stateManager.getCdcStateManager());
    }
  
    private List<AirbyteStateMessage> createDefaultState() {
      return List.of(new AirbyteStateMessage().withType(AirbyteStateType.STREAM).withStream(new AirbyteStreamState()));
    }
  
    private AirbyteStateMessage createStreamState(final String name,
                                                  final String namespace,
                                                  final List<String> cursorFields,
                                                  final String cursorValue) {
      final DbStreamState dbStreamState = new DbStreamState()
          .withStreamName(name)
          .withStreamNamespace(namespace);
  
      if (cursorFields != null && !cursorFields.isEmpty()) {
        dbStreamState.withCursorField(cursorFields);
      }
  
      if (cursorValue != null) {
        dbStreamState.withCursor(cursorValue);
      }
  
      return new AirbyteStateMessage()
          .withType(AirbyteStateType.STREAM)
          .withStream(new AirbyteStreamState()
              .withStreamDescriptor(new StreamDescriptor().withName(name).withNamespace(namespace))
              .withStreamState(Jsons.jsonNode(dbStreamState)));
    }
  
  }
  