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.beam.runners.dataflow.worker.streaming.sideinput;
  
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.contains;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertFalse;
  import static org.junit.Assert.assertNull;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.verifyNoMoreInteractions;
  import static org.mockito.Mockito.when;
  
  import java.io.Closeable;
  import java.util.Collections;
  import java.util.List;
  import java.util.concurrent.TimeUnit;
  import org.apache.beam.runners.dataflow.options.DataflowStreamingPipelineOptions;
  import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
  import org.apache.beam.runners.dataflow.worker.windmill.client.getdata.GetDataClient;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.coders.ListCoder;
  import org.apache.beam.sdk.coders.StringUtf8Coder;
  import org.apache.beam.sdk.coders.VarLongCoder;
  import org.apache.beam.sdk.coders.VoidCoder;
  import org.apache.beam.sdk.options.PipelineOptionsFactory;
  import org.apache.beam.sdk.testing.TestPipeline;
  import org.apache.beam.sdk.transforms.Create;
  import org.apache.beam.sdk.transforms.Sum;
  import org.apache.beam.sdk.transforms.View;
  import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
  import org.apache.beam.sdk.util.ByteStringOutputStream;
  import org.apache.beam.sdk.values.PCollectionView;
  import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
  import org.junit.Before;
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.rules.Timeout;
  import org.junit.runner.RunWith;
  import org.junit.runners.JUnit4;
  import org.mockito.ArgumentCaptor;
  import org.mockito.Mock;
  import org.mockito.MockitoAnnotations;
  
  /** Unit tests for {@link SideInputStateFetcher}. */
  // TODO: Add tests with different encoded windows to verify version is correctly plumbed.
  @SuppressWarnings("deprecation")
  @RunWith(JUnit4.class)
  public class SideInputStateFetcherTest {
    private static final String STATE_FAMILY = "state";
    @Rule public transient Timeout globalTimeout = Timeout.seconds(600);
    @Mock private GetDataClient server;
  
    @Mock private Supplier<Closeable> readStateSupplier;
  
    private static Windmill.GlobalData buildGlobalDataResponse(
        String tag, boolean isReady, ByteString data) {
      Windmill.GlobalData.Builder builder =
          Windmill.GlobalData.newBuilder()
              .setDataId(
                  Windmill.GlobalDataId.newBuilder()
                      .setTag(tag)
                      .setVersion(ByteString.EMPTY)
                      .build());
  
      if (isReady) {
        builder.setIsReady(true).setData(data);
      } else {
        builder.setIsReady(false);
      }
      return builder.build();
    }
  
    private static Windmill.GlobalDataRequest buildGlobalDataRequest(String tag, ByteString version) {
      Windmill.GlobalDataId id =
          Windmill.GlobalDataId.newBuilder().setTag(tag).setVersion(version).build();
  
      return Windmill.GlobalDataRequest.newBuilder()
          .setDataId(id)
          .setStateFamily(STATE_FAMILY)
          .setExistenceWatermarkDeadline(
              TimeUnit.MILLISECONDS.toMicros(GlobalWindow.INSTANCE.maxTimestamp().getMillis()))
          .build();
    }
  
    private static Windmill.GlobalDataRequest buildGlobalDataRequest(String tag) {
      return buildGlobalDataRequest(tag, ByteString.EMPTY);
    }
  
    @Before
    public void setUp() {
      MockitoAnnotations.initMocks(this);
    }
  
    @Test
    public void testFetchGlobalDataBasic() throws Exception {
      SideInputStateFetcherFactory factory =
          SideInputStateFetcherFactory.fromOptions(
              PipelineOptionsFactory.as(DataflowStreamingPipelineOptions.class));
      SideInputStateFetcher fetcher = factory.createSideInputStateFetcher(server::getSideInputData);
  
      ByteStringOutputStream stream = new ByteStringOutputStream();
      ListCoder.of(StringUtf8Coder.of())
          .encode(Collections.singletonList("data"), stream, Coder.Context.OUTER);
      ByteString encodedIterable = stream.toByteString();
  
      PCollectionView<String> view =
          TestPipeline.create().apply(Create.empty(StringUtf8Coder.of())).apply(View.asSingleton());
  
      String tag = view.getTagInternal().getId();
  
      // Test three calls in a row. First, data is not ready, then data is ready,
      // then the data is already cached.
      when(server.getSideInputData(any(Windmill.GlobalDataRequest.class)))
          .thenReturn(
              buildGlobalDataResponse(tag, false, null),
              buildGlobalDataResponse(tag, true, encodedIterable));
  
      assertFalse(
          fetcher
              .fetchSideInput(
                  view,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.UNKNOWN,
                  readStateSupplier)
              .isReady());
  
      assertFalse(
          fetcher
              .fetchSideInput(
                  view,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.UNKNOWN,
                  readStateSupplier)
              .isReady());
  
      assertEquals(
          "data",
          fetcher
              .fetchSideInput(
                  view,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.KNOWN_READY,
                  readStateSupplier)
              .value()
              .orElse(null));
      assertEquals(
          "data",
          fetcher
              .fetchSideInput(
                  view,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.KNOWN_READY,
                  readStateSupplier)
              .value()
              .orElse(null));
  
      verify(server, times(2)).getSideInputData(buildGlobalDataRequest(tag));
      verifyNoMoreInteractions(server);
    }
  
    @Test
    public void testFetchGlobalDataNull() throws Exception {
      SideInputStateFetcherFactory factory =
          SideInputStateFetcherFactory.fromOptions(
              PipelineOptionsFactory.as(DataflowStreamingPipelineOptions.class));
      SideInputStateFetcher fetcher = factory.createSideInputStateFetcher(server::getSideInputData);
  
      ByteStringOutputStream stream = new ByteStringOutputStream();
      ListCoder.of(VoidCoder.of())
          .encode(Collections.singletonList(null), stream, Coder.Context.OUTER);
      ByteString encodedIterable = stream.toByteString();
  
      PCollectionView<Void> view =
          TestPipeline.create().apply(Create.empty(VoidCoder.of())).apply(View.asSingleton());
  
      String tag = view.getTagInternal().getId();
  
      // Test three calls in a row. First, data is not ready, then data is ready,
      // then the data is already cached.
      when(server.getSideInputData(any(Windmill.GlobalDataRequest.class)))
          .thenReturn(
              buildGlobalDataResponse(tag, false, null),
              buildGlobalDataResponse(tag, true, encodedIterable));
  
      assertFalse(
          fetcher
              .fetchSideInput(
                  view,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.UNKNOWN,
                  readStateSupplier)
              .isReady());
      assertFalse(
          fetcher
              .fetchSideInput(
                  view,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.UNKNOWN,
                  readStateSupplier)
              .isReady());
      assertNull(
          fetcher
              .fetchSideInput(
                  view,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.KNOWN_READY,
                  readStateSupplier)
              .value()
              .orElse(null));
      assertNull(
          fetcher
              .fetchSideInput(
                  view,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.KNOWN_READY,
                  readStateSupplier)
              .value()
              .orElse(null));
  
      verify(server, times(2)).getSideInputData(buildGlobalDataRequest(tag));
      verifyNoMoreInteractions(server);
    }
  
    @Test
    public void testFetchGlobalDataCacheOverflow() throws Exception {
      Coder<List<String>> coder = ListCoder.of(StringUtf8Coder.of());
  
      ByteStringOutputStream stream = new ByteStringOutputStream();
      coder.encode(Collections.singletonList("data1"), stream, Coder.Context.OUTER);
      ByteString encodedIterable1 = stream.toByteStringAndReset();
      coder.encode(Collections.singletonList("data2"), stream, Coder.Context.OUTER);
      ByteString encodedIterable2 = stream.toByteString();
  
      Cache<SideInputCache.Key<?>, SideInput<?>> cache = CacheBuilder.newBuilder().build();
  
      SideInputStateFetcher fetcher =
          new SideInputStateFetcher(server::getSideInputData, new SideInputCache(cache));
  
      PCollectionView<String> view1 =
          TestPipeline.create().apply(Create.empty(StringUtf8Coder.of())).apply(View.asSingleton());
  
      PCollectionView<String> view2 =
          TestPipeline.create().apply(Create.empty(StringUtf8Coder.of())).apply(View.asSingleton());
  
      String tag1 = view1.getTagInternal().getId();
      String tag2 = view2.getTagInternal().getId();
  
      // Test four calls in a row. First, fetch view1, then view2 (which evicts view1 from the cache),
      // then view 1 again twice.
      when(server.getSideInputData(any(Windmill.GlobalDataRequest.class)))
          .thenReturn(
              buildGlobalDataResponse(tag1, true, encodedIterable1),
              buildGlobalDataResponse(tag2, true, encodedIterable2),
              buildGlobalDataResponse(tag1, true, encodedIterable1));
  
      assertEquals(
          "data1",
          fetcher
              .fetchSideInput(
                  view1,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.UNKNOWN,
                  readStateSupplier)
              .value()
              .orElse(null));
      assertEquals(
          "data2",
          fetcher
              .fetchSideInput(
                  view2,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.UNKNOWN,
                  readStateSupplier)
              .value()
              .orElse(null));
      cache.invalidateAll();
      assertEquals(
          "data1",
          fetcher
              .fetchSideInput(
                  view1,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.UNKNOWN,
                  readStateSupplier)
              .value()
              .orElse(null));
      assertEquals(
          "data1",
          fetcher
              .fetchSideInput(
                  view1,
                  GlobalWindow.INSTANCE,
                  STATE_FAMILY,
                  SideInputState.UNKNOWN,
                  readStateSupplier)
              .value()
              .orElse(null));
  
      ArgumentCaptor<Windmill.GlobalDataRequest> captor =
          ArgumentCaptor.forClass(Windmill.GlobalDataRequest.class);
  
      verify(server, times(3)).getSideInputData(captor.capture());
      verifyNoMoreInteractions(server);
  
      assertThat(
          captor.getAllValues(),
          contains(
              buildGlobalDataRequest(tag1),
              buildGlobalDataRequest(tag2),
              buildGlobalDataRequest(tag1)));
    }
  
    @Test
    public void testEmptyFetchGlobalData() {
      SideInputStateFetcherFactory factory =
          SideInputStateFetcherFactory.fromOptions(
              PipelineOptionsFactory.as(DataflowStreamingPipelineOptions.class));
      SideInputStateFetcher fetcher = factory.createSideInputStateFetcher(server::getSideInputData);
  
      ByteString encodedIterable = ByteString.EMPTY;
  
      PCollectionView<Long> view =
          TestPipeline.create()
              .apply(Create.empty(VarLongCoder.of()))
              .apply(Sum.longsGlobally().asSingletonView());
  
      String tag = view.getTagInternal().getId();
  
      // Test three calls in a row. First, data is not ready, then data is ready,
      // then the data is already cached.
      when(server.getSideInputData(any(Windmill.GlobalDataRequest.class)))
          .thenReturn(buildGlobalDataResponse(tag, true, encodedIterable));
  
      assertEquals(
          0L,
          (long)
              fetcher
                  .fetchSideInput(
                      view,
                      GlobalWindow.INSTANCE,
                      STATE_FAMILY,
                      SideInputState.UNKNOWN,
                      readStateSupplier)
                  .value()
                  .orElse(null));
  
      verify(server).getSideInputData(buildGlobalDataRequest(tag));
      verifyNoMoreInteractions(server);
    }
  }
  