diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/AddUuidsTransformTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/AddUuidsTransformTest.java deleted file mode 100644 index 331613e5d7f0..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/AddUuidsTransformTest.java +++ /dev/null @@ -1,97 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import com.google.cloud.pubsublite.proto.PubSubMessage; -import com.google.protobuf.ByteString; -import java.util.HashSet; -import java.util.Set; -import org.apache.beam.sdk.extensions.protobuf.ProtoCoder; -import org.apache.beam.sdk.testing.PAssert; -import org.apache.beam.sdk.testing.TestPipeline; -import org.apache.beam.sdk.testing.TestStream; -import org.apache.beam.sdk.transforms.SerializableFunction; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; -import org.joda.time.Duration; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public final class AddUuidsTransformTest { - @Rule public final TestPipeline pipeline = TestPipeline.create(); - - private static PubSubMessage newMessage(int identifier) { - return PubSubMessage.newBuilder() - .setKey(ByteString.copyFromUtf8(Integer.toString(identifier))) - .build(); - } - - private static SerializableFunction, Void> identifiersInAnyOrder( - Set identifiers) { - return messages -> { - Set uuids = new HashSet<>(); - messages.forEach( - message -> { - int identifier = Integer.parseInt(message.getKey().toStringUtf8()); - if (!identifiers.remove(identifier)) { - throw new IllegalStateException("Duplicate element " + identifier); - } - if (!uuids.add( - Uuid.of( - Iterables.getOnlyElement( - message.getAttributesMap().get(Uuid.DEFAULT_ATTRIBUTE).getValuesList())))) { - throw new IllegalStateException("Invalid duplicate Uuid: " + message.toString()); - } - }); - if (!identifiers.isEmpty()) { - throw new IllegalStateException("Elements not in collection: " + identifiers); - } - return null; - }; - } - - @Test - public void messagesSameBatch() { - TestStream messageStream = - TestStream.create(ProtoCoder.of(PubSubMessage.class)) - .addElements(newMessage(1), newMessage(2), newMessage(85)) - .advanceWatermarkToInfinity(); - PCollection outputs = - pipeline.apply(messageStream).apply(new AddUuidsTransform()); - PAssert.that(outputs).satisfies(identifiersInAnyOrder(Sets.newHashSet(1, 2, 85))); - pipeline.run(); - } - - @Test - public void messagesTimeDelayed() { - TestStream messageStream = - TestStream.create(ProtoCoder.of(PubSubMessage.class)) - .addElements(newMessage(1), newMessage(2)) - .advanceProcessingTime(Duration.standardDays(1)) - .addElements(newMessage(85)) - .advanceWatermarkToInfinity(); - PCollection outputs = - pipeline.apply(messageStream).apply(new AddUuidsTransform()); - PAssert.that(outputs).satisfies(identifiersInAnyOrder(Sets.newHashSet(1, 2, 85))); - pipeline.run(); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/BlockingCommitterImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/BlockingCommitterImplTest.java deleted file mode 100644 index 111c20c994b5..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/BlockingCommitterImplTest.java +++ /dev/null @@ -1,67 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.verify; - -import com.google.api.core.ApiFutures; -import com.google.cloud.pubsublite.Offset; -import com.google.cloud.pubsublite.internal.testing.FakeApiService; -import com.google.cloud.pubsublite.internal.wire.Committer; -import java.util.concurrent.TimeUnit; -import org.junit.Before; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.Spy; -import org.mockito.junit.MockitoJUnit; -import org.mockito.junit.MockitoRule; -import org.mockito.quality.Strictness; - -@RunWith(JUnit4.class) -public class BlockingCommitterImplTest { - - @Rule public MockitoRule mockito = MockitoJUnit.rule().strictness(Strictness.STRICT_STUBS); - - abstract static class FakeCommitter extends FakeApiService implements Committer {} - - @Spy private FakeCommitter fakeCommitter; - - private BlockingCommitter committer; - - @Before - public void setUp() { - fakeCommitter.startAsync().awaitRunning(); - committer = new BlockingCommitterImpl(fakeCommitter); - } - - @Test - public void commit() { - doReturn(ApiFutures.immediateFuture(null)).when(fakeCommitter).commitOffset(Offset.of(42)); - committer.commitOffset(Offset.of(42)); - } - - @Test - public void close() throws Exception { - committer.close(); - verify(fakeCommitter).stopAsync(); - verify(fakeCommitter).awaitTerminated(1, TimeUnit.MINUTES); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/CheckpointMarkImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/CheckpointMarkImplTest.java deleted file mode 100644 index fae47d0eec53..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/CheckpointMarkImplTest.java +++ /dev/null @@ -1,67 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import static org.junit.Assert.assertEquals; -import static org.mockito.Mockito.verify; - -import com.google.cloud.pubsublite.Offset; -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import org.junit.Before; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.Mock; -import org.mockito.junit.MockitoJUnit; -import org.mockito.junit.MockitoRule; -import org.mockito.quality.Strictness; - -@RunWith(JUnit4.class) -public class CheckpointMarkImplTest { - - private static final Offset OFFSET = Offset.of(42); - - @Rule public MockitoRule mockito = MockitoJUnit.rule().strictness(Strictness.STRICT_STUBS); - @Mock private BlockingCommitter committer; - - private CheckpointMarkImpl mark; - - @Before - public void setUp() { - mark = new CheckpointMarkImpl(OFFSET, () -> committer); - } - - @Test - public void testFinalize() throws Exception { - mark.finalizeCheckpoint(); - verify(committer).commitOffset(OFFSET); - } - - @Test - public void encodedCheckpointFinalizeFails() throws Exception { - ByteArrayOutputStream stream = new ByteArrayOutputStream(); - CheckpointMarkImpl.coder().encode(mark, stream); - CheckpointMarkImpl impl = - CheckpointMarkImpl.coder().decode(new ByteArrayInputStream(stream.toByteArray())); - assertEquals(impl.offset, OFFSET); - // No exception thrown despite throw to work around DirectRunner issue. - impl.finalizeCheckpoint(); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/FakeSerializable.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/FakeSerializable.java deleted file mode 100644 index ded551bb9b05..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/FakeSerializable.java +++ /dev/null @@ -1,59 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import java.io.Serializable; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.atomic.AtomicInteger; -import org.apache.beam.sdk.util.SerializableSupplier; - -/** - * A FakeSerializable hides a non-serializable object in a static map and returns a handle into the - * static map. It is useful in the presence of in-process serialization, but not out of process - * serialization. - */ -final class FakeSerializable { - private static final AtomicInteger idCounter = new AtomicInteger(0); - private static final ConcurrentHashMap map = new ConcurrentHashMap<>(); - - private FakeSerializable() {} - - static class Handle implements Serializable { - private Handle(int id) { - this.id = id; - } - - private final int id; - - @SuppressWarnings("unchecked") - T get() { - return (T) map.get(id); - } - } - - static Handle put(T value) { - int id = idCounter.incrementAndGet(); - map.put(id, value); - return new Handle(id); - } - - static SerializableSupplier getSupplier(T value) { - Handle handle = put(value); - return handle::get; - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriberImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriberImplTest.java deleted file mode 100644 index 602672704078..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/MemoryBufferedSubscriberImplTest.java +++ /dev/null @@ -1,180 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import static com.google.cloud.pubsublite.internal.testing.UnitTestExamples.example; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertThrows; -import static org.junit.Assert.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyLong; -import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.reset; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.MockitoAnnotations.initMocks; - -import com.google.cloud.pubsublite.Offset; -import com.google.cloud.pubsublite.Partition; -import com.google.cloud.pubsublite.internal.testing.FakeApiService; -import com.google.cloud.pubsublite.internal.wire.Subscriber; -import com.google.cloud.pubsublite.proto.Cursor; -import com.google.cloud.pubsublite.proto.SequencedMessage; -import java.util.List; -import java.util.function.Consumer; -import java.util.function.Function; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; -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.Mock; -import org.mockito.Spy; - -@RunWith(JUnit4.class) -public class MemoryBufferedSubscriberImplTest { - private static final long MAX_MEMORY = 1024; - - @Rule public Timeout globalTimeout = Timeout.seconds(30); - - abstract static class FakeSubscriber extends FakeApiService implements Subscriber {} - - @Spy FakeSubscriber subscriber; - @Mock Function>, Subscriber> subscriberFactory; - @Mock MemoryLimiter limiter; - @Mock MemoryLimiter.Block block; - - MemoryBufferedSubscriber bufferedSubscriber; - - Consumer> consumer; - - private static SequencedMessage messageWithSize(long size) { - return SequencedMessage.newBuilder().setSizeBytes(size).build(); - } - - @Before - public void setUp() { - initMocks(this); - doAnswer( - args -> { - consumer = args.getArgument(0); - return subscriber; - }) - .when(subscriberFactory) - .apply(any()); - doReturn(1L).when(limiter).minBlockSize(); - doReturn(MAX_MEMORY).when(limiter).maxBlockSize(); - checkNotNull(block); - checkNotNull(limiter); - doReturn(block).when(limiter).claim(anyLong()); - doReturn(MAX_MEMORY).when(block).claimed(); - bufferedSubscriber = - new MemoryBufferedSubscriberImpl( - example(Partition.class), example(Offset.class), limiter, subscriberFactory); - checkNotNull(consumer); - bufferedSubscriber.startAsync().awaitRunning(); - verify(subscriber).startAsync(); - assertTrue(subscriber.isRunning()); - } - - @Test - public void underlyingFailureFails() { - subscriber.fail(new RuntimeException("bad")); - assertThrows(Exception.class, subscriber::awaitTerminated); - } - - @Test - public void rebufferReducesToOutstandingWhenLittleData() { - consumer.accept(ImmutableList.of(messageWithSize(MAX_MEMORY / 4))); - bufferedSubscriber.pop(); - bufferedSubscriber.rebuffer(); - verify(block).close(); - verify(limiter).claim(3 * MAX_MEMORY / 4); - } - - @Test - public void rebufferCannotGoBelowMin() { - long minBlock = MAX_MEMORY * 4 / 5; - doReturn(minBlock).when(limiter).minBlockSize(); - for (int i = 0; i < 1000; ++i) { - // Rebuffer many times with no data to bring down the target value - bufferedSubscriber.rebuffer(); - } - reset(limiter); - doReturn(minBlock).when(limiter).minBlockSize(); - doReturn(block).when(limiter).claim(anyLong()); - // Deliver enough data that 3 * minBlock / 4 is outstanding, buffer is allowed to and will - // shrink except it is limited by min block size. - consumer.accept(ImmutableList.of(messageWithSize(2 * MAX_MEMORY / 5))); - bufferedSubscriber.pop(); - bufferedSubscriber.rebuffer(); - verify(limiter).claim(minBlock); - } - - @Test - public void rebufferStaysSameOnHalfDelivered() { - consumer.accept(ImmutableList.of(messageWithSize(MAX_MEMORY / 4))); - bufferedSubscriber.pop(); - bufferedSubscriber.rebuffer(); - verify(limiter).claim(3 * MAX_MEMORY / 4); - consumer.accept(ImmutableList.of(messageWithSize(3 * MAX_MEMORY / 8))); - bufferedSubscriber.pop(); - bufferedSubscriber.rebuffer(); - verify(limiter).claim(3 * MAX_MEMORY / 4); - } - - @Test - public void rebufferGrowsOnMoreDelivered() { - consumer.accept(ImmutableList.of(messageWithSize(MAX_MEMORY / 4))); - bufferedSubscriber.pop(); - bufferedSubscriber.rebuffer(); - verify(limiter).claim(3 * MAX_MEMORY / 4); - consumer.accept( - ImmutableList.of(messageWithSize(MAX_MEMORY / 2), messageWithSize(MAX_MEMORY / 8))); - bufferedSubscriber.rebuffer(); - verify(limiter, times(2)).claim(MAX_MEMORY); // once in setup - } - - @Test - public void dataAvailableToCaller() { - SequencedMessage message1 = - SequencedMessage.newBuilder() - .setCursor(Cursor.newBuilder().setOffset(example(Offset.class).value() + 10)) - .setSizeBytes(1) - .build(); - SequencedMessage message2 = - SequencedMessage.newBuilder() - .setCursor(Cursor.newBuilder().setOffset(example(Offset.class).value() + 20)) - .setSizeBytes(1) - .build(); - assertFalse(bufferedSubscriber.peek().isPresent()); - consumer.accept(ImmutableList.of(message1, message2)); - assertEquals(bufferedSubscriber.fetchOffset(), example(Offset.class)); - assertEquals(bufferedSubscriber.peek().get(), message1); - bufferedSubscriber.pop(); - assertEquals(bufferedSubscriber.fetchOffset(), Offset.of(message1.getCursor().getOffset() + 1)); - assertEquals(bufferedSubscriber.peek().get(), message2); - bufferedSubscriber.pop(); - assertEquals(bufferedSubscriber.fetchOffset(), Offset.of(message2.getCursor().getOffset() + 1)); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTrackerTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTrackerTest.java deleted file mode 100644 index 667fd52a8b1f..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/OffsetByteRangeTrackerTest.java +++ /dev/null @@ -1,135 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertNotNull; -import static org.junit.Assert.assertNull; -import static org.junit.Assert.assertThrows; -import static org.junit.Assert.assertTrue; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; -import static org.mockito.MockitoAnnotations.initMocks; - -import com.google.api.gax.rpc.ApiException; -import com.google.api.gax.rpc.StatusCode.Code; -import com.google.cloud.pubsublite.Offset; -import com.google.cloud.pubsublite.internal.CheckedApiException; -import com.google.cloud.pubsublite.proto.ComputeMessageStatsResponse; -import org.apache.beam.sdk.io.range.OffsetRange; -import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress; -import org.apache.beam.sdk.transforms.splittabledofn.SplitResult; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Ticker; -import org.junit.Before; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.Spy; - -@RunWith(JUnit4.class) -@SuppressWarnings("initialization.fields.uninitialized") -public class OffsetByteRangeTrackerTest { - private static final double IGNORED_FRACTION = -10000000.0; - private static final long MIN_BYTES = 1000; - private static final OffsetRange RANGE = new OffsetRange(123L, Long.MAX_VALUE); - private final TopicBacklogReader unownedBacklogReader = mock(TopicBacklogReader.class); - - @Spy Ticker ticker; - private OffsetByteRangeTracker tracker; - - @Before - public void setUp() { - initMocks(this); - when(ticker.read()).thenReturn(0L); - tracker = new OffsetByteRangeTracker(OffsetByteRange.of(RANGE, 0), unownedBacklogReader); - } - - @Test - public void progressTracked() { - assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(123), 10))); - assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(124), 11))); - when(unownedBacklogReader.computeMessageStats(Offset.of(125))) - .thenReturn(ComputeMessageStatsResponse.newBuilder().setMessageBytes(1000).build()); - Progress progress = tracker.getProgress(); - assertEquals(21, progress.getWorkCompleted(), .0001); - assertEquals(1000, progress.getWorkRemaining(), .0001); - } - - @Test - public void getProgressStatsFailure() { - when(unownedBacklogReader.computeMessageStats(Offset.of(123))) - .thenThrow(new CheckedApiException(Code.INTERNAL).underlying); - assertThrows(ApiException.class, tracker::getProgress); - } - - @Test - @SuppressWarnings({"dereference.of.nullable", "argument"}) - public void claimSplitSuccess() { - assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_000), MIN_BYTES))); - assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(10_000), MIN_BYTES))); - SplitResult splits = tracker.trySplit(IGNORED_FRACTION); - OffsetByteRange primary = splits.getPrimary(); - assertEquals(RANGE.getFrom(), primary.getRange().getFrom()); - assertEquals(10_001, primary.getRange().getTo()); - assertEquals(MIN_BYTES * 2, primary.getByteCount()); - OffsetByteRange residual = splits.getResidual(); - assertEquals(10_001, residual.getRange().getFrom()); - assertEquals(Long.MAX_VALUE, residual.getRange().getTo()); - assertEquals(0, residual.getByteCount()); - assertEquals(splits.getPrimary(), tracker.currentRestriction()); - tracker.checkDone(); - assertNull(tracker.trySplit(IGNORED_FRACTION)); - } - - @Test - @SuppressWarnings({"dereference.of.nullable", "argument"}) - public void splitWithoutClaimEmpty() { - when(ticker.read()).thenReturn(100000000000000L); - SplitResult splits = tracker.trySplit(IGNORED_FRACTION); - assertEquals(RANGE.getFrom(), splits.getPrimary().getRange().getFrom()); - assertEquals(RANGE.getFrom(), splits.getPrimary().getRange().getTo()); - assertEquals(RANGE, splits.getResidual().getRange()); - assertEquals(splits.getPrimary(), tracker.currentRestriction()); - tracker.checkDone(); - assertNull(tracker.trySplit(IGNORED_FRACTION)); - } - - @Test - public void unboundedNotDone() { - assertThrows(IllegalStateException.class, tracker::checkDone); - } - - @Test - public void cannotClaimBackwards() { - assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_000), MIN_BYTES))); - assertThrows( - IllegalArgumentException.class, - () -> tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_000), MIN_BYTES))); - assertThrows( - IllegalArgumentException.class, - () -> tracker.tryClaim(OffsetByteProgress.of(Offset.of(999), MIN_BYTES))); - } - - @Test - public void cannotClaimSplitRange() { - assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_000), MIN_BYTES))); - assertNotNull(tracker.trySplit(IGNORED_FRACTION)); - assertFalse(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_001), MIN_BYTES))); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerSubscriptionPartitionSdfTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerSubscriptionPartitionSdfTest.java deleted file mode 100644 index 7047d35b21ac..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PerSubscriptionPartitionSdfTest.java +++ /dev/null @@ -1,205 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import static com.google.cloud.pubsublite.internal.testing.UnitTestExamples.example; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertSame; -import static org.junit.Assert.assertThrows; -import static org.junit.Assert.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.inOrder; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.verifyNoInteractions; -import static org.mockito.Mockito.when; -import static org.mockito.MockitoAnnotations.initMocks; - -import com.google.api.gax.rpc.ApiException; -import com.google.api.gax.rpc.StatusCode.Code; -import com.google.cloud.pubsublite.Offset; -import com.google.cloud.pubsublite.Partition; -import com.google.cloud.pubsublite.SubscriptionPath; -import com.google.cloud.pubsublite.internal.CheckedApiException; -import com.google.cloud.pubsublite.proto.SequencedMessage; -import java.io.ByteArrayOutputStream; -import java.io.ObjectOutputStream; -import java.util.Optional; -import javax.annotation.Nonnull; -import org.apache.beam.sdk.io.range.OffsetRange; -import org.apache.beam.sdk.transforms.DoFn.OutputReceiver; -import org.apache.beam.sdk.transforms.DoFn.ProcessContinuation; -import org.apache.beam.sdk.transforms.SerializableBiFunction; -import org.apache.beam.sdk.transforms.SerializableFunction; -import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; -import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.math.DoubleMath; -import org.junit.Before; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.InOrder; -import org.mockito.Mock; -import org.mockito.Spy; - -@RunWith(JUnit4.class) -@SuppressWarnings("initialization.fields.uninitialized") -public class PerSubscriptionPartitionSdfTest { - - private static final OffsetByteRange RESTRICTION = - OffsetByteRange.of(new OffsetRange(1, Long.MAX_VALUE), 0); - private static final SubscriptionPartition PARTITION = - SubscriptionPartition.of(example(SubscriptionPath.class), example(Partition.class)); - - @Mock SerializableFunction offsetReaderFactory; - - @Mock ManagedFactory backlogReaderFactory; - @Mock TopicBacklogReader backlogReader; - - @Mock - SerializableBiFunction trackerFactory; - - @Mock SubscriptionPartitionProcessorFactory processorFactory; - @Mock ManagedFactory committerFactory; - - @Mock InitialOffsetReader initialOffsetReader; - @Spy TrackerWithProgress tracker; - @Mock OutputReceiver output; - @Mock SubscriptionPartitionProcessor processor; - @Mock BlockingCommitter committer; - - PerSubscriptionPartitionSdf sdf; - - @Before - public void setUp() { - initMocks(this); - when(offsetReaderFactory.apply(any())).thenReturn(initialOffsetReader); - when(processorFactory.newProcessor(any(), any(), any())).thenReturn(processor); - when(trackerFactory.apply(any(), any())).thenReturn(tracker); - when(committerFactory.create(any())).thenReturn(committer); - when(tracker.currentRestriction()).thenReturn(RESTRICTION); - when(backlogReaderFactory.create(any())).thenReturn(backlogReader); - sdf = - new PerSubscriptionPartitionSdf( - backlogReaderFactory, - committerFactory, - offsetReaderFactory, - trackerFactory, - processorFactory); - } - - @Test - public void getInitialRestrictionReadSuccess() { - when(initialOffsetReader.read()).thenReturn(example(Offset.class)); - OffsetByteRange range = sdf.getInitialRestriction(PARTITION); - assertEquals(example(Offset.class).value(), range.getRange().getFrom()); - assertEquals(Long.MAX_VALUE, range.getRange().getTo()); - assertEquals(0, range.getByteCount()); - verify(offsetReaderFactory).apply(PARTITION); - } - - @Test - public void getInitialRestrictionReadFailure() { - when(initialOffsetReader.read()).thenThrow(new CheckedApiException(Code.INTERNAL).underlying); - assertThrows(ApiException.class, () -> sdf.getInitialRestriction(PARTITION)); - } - - @Test - public void newTrackerCallsFactory() { - assertSame(tracker, sdf.newTracker(PARTITION, RESTRICTION)); - verify(trackerFactory).apply(backlogReader, RESTRICTION); - } - - @Test - public void tearDownClosesBacklogReaderFactory() throws Exception { - sdf.teardown(); - verify(backlogReaderFactory).close(); - } - - @Test - @SuppressWarnings("argument") - public void process() throws Exception { - when(processor.run()).thenReturn(ProcessContinuation.resume()); - when(processorFactory.newProcessor(any(), any(), any())) - .thenAnswer( - args -> { - @Nonnull - RestrictionTracker wrapped = args.getArgument(1); - when(tracker.tryClaim(any())).thenReturn(true).thenReturn(false); - assertTrue(wrapped.tryClaim(OffsetByteProgress.of(example(Offset.class), 123))); - assertFalse(wrapped.tryClaim(OffsetByteProgress.of(Offset.of(333333), 123))); - return processor; - }); - doReturn(Optional.of(example(Offset.class))).when(processor).lastClaimed(); - assertEquals(ProcessContinuation.resume(), sdf.processElement(tracker, PARTITION, output)); - verify(processorFactory).newProcessor(eq(PARTITION), any(), eq(output)); - InOrder order = inOrder(processor); - order.verify(processor).run(); - order.verify(processor).lastClaimed(); - InOrder order2 = inOrder(committerFactory, committer); - order2.verify(committerFactory).create(PARTITION); - order2.verify(committer).commitOffset(Offset.of(example(Offset.class).value() + 1)); - } - - private static final class NoopManagedFactory - implements ManagedFactory { - - @Override - public T create(SubscriptionPartition subscriptionPartition) { - return null; - } - - @Override - public void close() {} - } - - @Test - @SuppressWarnings("return") - public void dofnIsSerializable() throws Exception { - ObjectOutputStream output = new ObjectOutputStream(new ByteArrayOutputStream()); - output.writeObject( - new PerSubscriptionPartitionSdf( - new NoopManagedFactory<>(), - new NoopManagedFactory<>(), - (x) -> null, - (x, y) -> null, - (x, y, z) -> null)); - } - - @Test - public void getProgressUnboundedRangeDelegates() { - Progress progress = Progress.from(0, 0.2); - when(tracker.getProgress()).thenReturn(progress); - assertTrue( - DoubleMath.fuzzyEquals( - progress.getWorkRemaining(), sdf.getSize(PARTITION, RESTRICTION), .0001)); - verify(tracker).getProgress(); - } - - @Test - public void getProgressBoundedReturnsBytes() { - assertTrue( - DoubleMath.fuzzyEquals( - 123.0, - sdf.getSize(PARTITION, OffsetByteRange.of(new OffsetRange(87, 8000), 123)), - .0001)); - verifyNoInteractions(tracker); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PubsubLiteSinkTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PubsubLiteSinkTest.java deleted file mode 100644 index ee277b8f8ce8..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/PubsubLiteSinkTest.java +++ /dev/null @@ -1,195 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.containsInAnyOrder; -import static org.hamcrest.Matchers.equalTo; -import static org.junit.Assert.assertThrows; -import static org.junit.Assert.assertTrue; -import static org.junit.Assert.fail; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import com.google.api.core.ApiFutures; -import com.google.api.core.SettableApiFuture; -import com.google.api.gax.rpc.StatusCode.Code; -import com.google.cloud.pubsublite.CloudRegion; -import com.google.cloud.pubsublite.CloudZone; -import com.google.cloud.pubsublite.MessageMetadata; -import com.google.cloud.pubsublite.Offset; -import com.google.cloud.pubsublite.Partition; -import com.google.cloud.pubsublite.ProjectNumber; -import com.google.cloud.pubsublite.TopicName; -import com.google.cloud.pubsublite.TopicPath; -import com.google.cloud.pubsublite.internal.CheckedApiException; -import com.google.cloud.pubsublite.internal.ExtractStatus; -import com.google.cloud.pubsublite.internal.Publisher; -import com.google.cloud.pubsublite.internal.testing.FakeApiService; -import com.google.cloud.pubsublite.proto.PubSubMessage; -import com.google.protobuf.ByteString; -import java.util.Arrays; -import java.util.Optional; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.stream.Collectors; -import org.apache.beam.sdk.Pipeline.PipelineExecutionException; -import org.apache.beam.sdk.io.gcp.pubsublite.PublisherOptions; -import org.apache.beam.sdk.testing.TestPipeline; -import org.apache.beam.sdk.transforms.Create; -import org.apache.beam.sdk.transforms.ParDo; -import org.junit.Before; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.ArgumentCaptor; -import org.mockito.Captor; -import org.mockito.MockitoAnnotations; -import org.mockito.Spy; - -@RunWith(JUnit4.class) -public class PubsubLiteSinkTest { - @Rule public final TestPipeline pipeline = TestPipeline.create(); - - abstract static class PublisherFakeService extends FakeApiService - implements Publisher {} - - @Spy private PublisherFakeService publisher; - - private PublisherOptions defaultOptions() { - return PublisherOptions.newBuilder() - .setTopicPath( - TopicPath.newBuilder() - .setProject(ProjectNumber.of(9)) - .setName(TopicName.of("abc")) - .setLocation(CloudZone.of(CloudRegion.of("us-east1"), 'a')) - .build()) - .build(); - } - - private final PubsubLiteSink sink = new PubsubLiteSink(defaultOptions()); - - @Captor - final ArgumentCaptor publishedMessageCaptor = - ArgumentCaptor.forClass(PubSubMessage.class); - - private void runWith(PubSubMessage... messages) { - pipeline - .apply(Create.of(Arrays.stream(messages).collect(Collectors.toList()))) - .apply(ParDo.of(sink)); - pipeline.run(); - } - - @Before - public void setUp() throws Exception { - MockitoAnnotations.initMocks(this); - PerServerPublisherCache.PUBLISHER_CACHE.set(defaultOptions(), publisher); - } - - @Test - public void singleMessagePublishes() throws Exception { - when(publisher.publish(PubSubMessage.newBuilder().build())) - .thenReturn(ApiFutures.immediateFuture(MessageMetadata.of(Partition.of(1), Offset.of(2)))); - runWith(PubSubMessage.newBuilder().build()); - verify(publisher).publish(PubSubMessage.newBuilder().build()); - } - - @Test - public void manyMessagePublishes() throws Exception { - PubSubMessage message1 = PubSubMessage.newBuilder().build(); - PubSubMessage message2 = - PubSubMessage.newBuilder().setKey(ByteString.copyFromUtf8("abc")).build(); - when(publisher.publish(message1)) - .thenReturn(ApiFutures.immediateFuture(MessageMetadata.of(Partition.of(1), Offset.of(2)))); - when(publisher.publish(message2)) - .thenReturn(ApiFutures.immediateFuture(MessageMetadata.of(Partition.of(85), Offset.of(3)))); - runWith(message1, message2); - verify(publisher, times(2)).publish(publishedMessageCaptor.capture()); - assertThat(publishedMessageCaptor.getAllValues(), containsInAnyOrder(message1, message2)); - } - - @Test - public void singleExceptionWhenProcessing() { - PubSubMessage message1 = PubSubMessage.newBuilder().build(); - when(publisher.publish(message1)) - .thenReturn( - ApiFutures.immediateFailedFuture(new CheckedApiException(Code.INTERNAL).underlying)); - PipelineExecutionException e = - assertThrows(PipelineExecutionException.class, () -> runWith(message1)); - verify(publisher).publish(message1); - Optional statusOr = ExtractStatus.extract(e.getCause()); - assertTrue(statusOr.isPresent()); - assertThat(statusOr.get().code(), equalTo(Code.INTERNAL)); - } - - @Test - public void exceptionMixedWithOK() throws Exception { - PubSubMessage message1 = PubSubMessage.newBuilder().build(); - PubSubMessage message2 = - PubSubMessage.newBuilder().setKey(ByteString.copyFromUtf8("abc")).build(); - PubSubMessage message3 = - PubSubMessage.newBuilder().setKey(ByteString.copyFromUtf8("def")).build(); - SettableApiFuture future1 = SettableApiFuture.create(); - SettableApiFuture future2 = SettableApiFuture.create(); - SettableApiFuture future3 = SettableApiFuture.create(); - CountDownLatch startedLatch = new CountDownLatch(3); - when(publisher.publish(message1)) - .then( - invocation -> { - startedLatch.countDown(); - return future1; - }); - when(publisher.publish(message2)) - .then( - invocation -> { - startedLatch.countDown(); - return future2; - }); - when(publisher.publish(message3)) - .then( - invocation -> { - startedLatch.countDown(); - return future3; - }); - ExecutorService exec = Executors.newCachedThreadPool(); - exec.execute( - () -> { - try { - startedLatch.await(); - future1.set(MessageMetadata.of(Partition.of(1), Offset.of(2))); - future2.setException(new CheckedApiException(Code.INTERNAL).underlying); - future3.set(MessageMetadata.of(Partition.of(1), Offset.of(3))); - } catch (InterruptedException e) { - fail(); - throw new RuntimeException(e); - } - }); - PipelineExecutionException e = - assertThrows(PipelineExecutionException.class, () -> runWith(message1, message2, message3)); - verify(publisher, times(3)).publish(publishedMessageCaptor.capture()); - assertThat( - publishedMessageCaptor.getAllValues(), containsInAnyOrder(message1, message2, message3)); - Optional statusOr = ExtractStatus.extract(e.getCause()); - assertTrue(statusOr.isPresent()); - assertThat(statusOr.get().code(), equalTo(Code.INTERNAL)); - exec.shutdownNow(); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionLoaderTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionLoaderTest.java deleted file mode 100644 index 5d4d99beaeab..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionLoaderTest.java +++ /dev/null @@ -1,91 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import static com.google.cloud.pubsublite.internal.testing.UnitTestExamples.example; -import static org.mockito.Mockito.when; -import static org.mockito.MockitoAnnotations.initMocks; - -import com.google.cloud.pubsublite.Partition; -import com.google.cloud.pubsublite.SubscriptionPath; -import com.google.cloud.pubsublite.TopicPath; -import org.apache.beam.sdk.testing.PAssert; -import org.apache.beam.sdk.testing.TestPipeline; -import org.apache.beam.sdk.transforms.SerializableFunction; -import org.apache.beam.sdk.util.SerializableSupplier; -import org.apache.beam.sdk.values.PCollection; -import org.joda.time.Duration; -import org.junit.Before; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.Mock; - -@SuppressWarnings("initialization.fields.uninitialized") -@RunWith(JUnit4.class) -public class SubscriptionPartitionLoaderTest { - @Rule public final transient TestPipeline pipeline = TestPipeline.create(); - @Mock SerializableFunction getPartitionCount; - - @Mock SerializableSupplier terminate; - private SubscriptionPartitionLoader loader; - - @Before - public void setUp() { - initMocks(this); - FakeSerializable.Handle> handle = - FakeSerializable.put(getPartitionCount); - FakeSerializable.Handle> terminateHandle = - FakeSerializable.put(terminate); - loader = - new SubscriptionPartitionLoader( - example(TopicPath.class), - example(SubscriptionPath.class), - topic -> handle.get().apply(topic), - Duration.millis(50), - () -> terminateHandle.get().get()); - } - - @Test - public void singleResult() { - when(getPartitionCount.apply(example(TopicPath.class))).thenReturn(3); - when(terminate.get()).thenReturn(false).thenReturn(false).thenReturn(true); - PCollection output = pipeline.apply(loader); - PAssert.that(output) - .containsInAnyOrder( - SubscriptionPartition.of(example(SubscriptionPath.class), Partition.of(0)), - SubscriptionPartition.of(example(SubscriptionPath.class), Partition.of(1)), - SubscriptionPartition.of(example(SubscriptionPath.class), Partition.of(2))); - pipeline.run().waitUntilFinish(); - } - - @Test - public void addedResults() { - when(getPartitionCount.apply(example(TopicPath.class))).thenReturn(3).thenReturn(4); - when(terminate.get()).thenReturn(false).thenReturn(false).thenReturn(true); - PCollection output = pipeline.apply(loader); - PAssert.that(output) - .containsInAnyOrder( - SubscriptionPartition.of(example(SubscriptionPath.class), Partition.of(0)), - SubscriptionPartition.of(example(SubscriptionPath.class), Partition.of(1)), - SubscriptionPartition.of(example(SubscriptionPath.class), Partition.of(2)), - SubscriptionPartition.of(example(SubscriptionPath.class), Partition.of(3))); - pipeline.run().waitUntilFinish(); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImplTest.java deleted file mode 100644 index a18580fe5c6f..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/SubscriptionPartitionProcessorImplTest.java +++ /dev/null @@ -1,146 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import static com.google.cloud.pubsublite.internal.testing.UnitTestExamples.example; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.inOrder; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; -import static org.mockito.MockitoAnnotations.initMocks; - -import com.google.api.gax.rpc.ApiException; -import com.google.api.gax.rpc.StatusCode.Code; -import com.google.cloud.pubsublite.Offset; -import com.google.cloud.pubsublite.internal.CheckedApiException; -import com.google.cloud.pubsublite.internal.testing.FakeApiService; -import com.google.cloud.pubsublite.proto.Cursor; -import com.google.cloud.pubsublite.proto.SequencedMessage; -import com.google.protobuf.util.Timestamps; -import java.util.Optional; -import org.apache.beam.sdk.io.range.OffsetRange; -import org.apache.beam.sdk.transforms.DoFn.OutputReceiver; -import org.apache.beam.sdk.transforms.DoFn.ProcessContinuation; -import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; -import org.joda.time.Instant; -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.InOrder; -import org.mockito.Mock; -import org.mockito.Spy; - -@RunWith(JUnit4.class) -@SuppressWarnings("initialization.fields.uninitialized") -public class SubscriptionPartitionProcessorImplTest { - @Spy RestrictionTracker tracker; - @Mock OutputReceiver receiver; - - @Rule public Timeout globalTimeout = Timeout.seconds(30); - - abstract static class FakeSubscriber extends FakeApiService implements MemoryBufferedSubscriber {} - - @Spy FakeSubscriber subscriber; - - private static SequencedMessage messageWithOffset(long offset) { - return SequencedMessage.newBuilder() - .setCursor(Cursor.newBuilder().setOffset(offset)) - .setPublishTime(Timestamps.fromMillis(10000 + offset)) - .setSizeBytes(1024) - .build(); - } - - private OffsetByteRange initialRange() { - return OffsetByteRange.of(new OffsetRange(example(Offset.class).value(), Long.MAX_VALUE)); - } - - @Before - public void setUp() { - initMocks(this); - subscriber.startAsync().awaitRunning(); - when(tracker.currentRestriction()).thenReturn(initialRange()); - doReturn(example(Offset.class)).when(subscriber).fetchOffset(); - } - - private SubscriptionPartitionProcessor newProcessor() { - return new SubscriptionPartitionProcessorImpl(tracker, receiver, subscriber); - } - - @Test - public void create() { - SubscriptionPartitionProcessor processor = newProcessor(); - assertEquals(ProcessContinuation.resume(), processor.run()); - InOrder order = inOrder(subscriber); - order.verify(subscriber).fetchOffset(); - order.verify(subscriber).rebuffer(); - } - - @Test - public void createRebufferThrows() throws Exception { - doThrow(new CheckedApiException(Code.OUT_OF_RANGE).underlying).when(subscriber).rebuffer(); - assertThrows(ApiException.class, this::newProcessor); - } - - @Test - public void failedClaimCausesStop() { - SubscriptionPartitionProcessor processor = newProcessor(); - - when(tracker.tryClaim(any())).thenReturn(false); - doReturn(Optional.of(messageWithOffset(1))).when(subscriber).peek(); - - assertEquals(ProcessContinuation.stop(), processor.run()); - - verify(tracker, times(1)).tryClaim(any()); - verify(subscriber, times(0)).pop(); - assertFalse(processor.lastClaimed().isPresent()); - } - - @Test - public void successfulClaimsThenNoMoreMessagesFromSubscriber() { - doReturn(true).when(tracker).tryClaim(any()); - - SequencedMessage message1 = messageWithOffset(1); - SequencedMessage message3 = messageWithOffset(3); - doReturn(Optional.of(message1), Optional.of(message3), Optional.empty()) - .when(subscriber) - .peek(); - - SubscriptionPartitionProcessor processor = newProcessor(); - assertEquals(ProcessContinuation.resume(), processor.run()); - - InOrder order = inOrder(tracker, receiver); - order.verify(tracker).tryClaim(OffsetByteProgress.of(Offset.of(1), message1.getSizeBytes())); - order - .verify(receiver) - .outputWithTimestamp(message1, new Instant(Timestamps.toMillis(message1.getPublishTime()))); - order.verify(tracker).tryClaim(OffsetByteProgress.of(Offset.of(3), message3.getSizeBytes())); - order - .verify(receiver) - .outputWithTimestamp(message3, new Instant(Timestamps.toMillis(message3.getPublishTime()))); - assertEquals(processor.lastClaimed().get(), Offset.of(3)); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/TopicBacklogReaderImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/TopicBacklogReaderImplTest.java deleted file mode 100644 index e60e8c0f9fbf..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/TopicBacklogReaderImplTest.java +++ /dev/null @@ -1,99 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import static com.google.cloud.pubsublite.internal.testing.UnitTestExamples.example; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertThrows; -import static org.mockito.Mockito.when; -import static org.mockito.MockitoAnnotations.initMocks; - -import com.google.api.core.ApiFutures; -import com.google.api.gax.rpc.ApiException; -import com.google.api.gax.rpc.StatusCode.Code; -import com.google.cloud.pubsublite.Offset; -import com.google.cloud.pubsublite.Partition; -import com.google.cloud.pubsublite.TopicPath; -import com.google.cloud.pubsublite.internal.CheckedApiException; -import com.google.cloud.pubsublite.internal.TopicStatsClient; -import com.google.cloud.pubsublite.proto.ComputeMessageStatsResponse; -import com.google.protobuf.Timestamp; -import org.junit.Before; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.Mock; -import org.mockito.junit.MockitoJUnit; -import org.mockito.junit.MockitoRule; - -@SuppressWarnings("uninitialized") -@RunWith(JUnit4.class) -public final class TopicBacklogReaderImplTest { - - @Rule public final MockitoRule mockito = MockitoJUnit.rule(); - - @Mock TopicStatsClient mockClient; - - private TopicBacklogReader reader; - - @Before - public void setUp() { - initMocks(this); - this.reader = - new TopicBacklogReaderImpl(mockClient, example(TopicPath.class), example(Partition.class)); - } - - @SuppressWarnings("incompatible") - @Test - public void computeMessageStats_failure() { - when(mockClient.computeMessageStats( - example(TopicPath.class), - example(Partition.class), - example(Offset.class), - Offset.of(Long.MAX_VALUE))) - .thenReturn( - ApiFutures.immediateFailedFuture(new CheckedApiException(Code.UNAVAILABLE).underlying)); - - ApiException e = - assertThrows(ApiException.class, () -> reader.computeMessageStats(example(Offset.class))); - assertEquals(Code.UNAVAILABLE, e.getStatusCode().getCode()); - } - - @Test - public void computeMessageStats_validResponseCached() { - Timestamp minEventTime = Timestamp.newBuilder().setSeconds(1000).setNanos(10).build(); - Timestamp minPublishTime = Timestamp.newBuilder().setSeconds(1001).setNanos(11).build(); - ComputeMessageStatsResponse response = - ComputeMessageStatsResponse.newBuilder() - .setMessageCount(10) - .setMessageBytes(100) - .setMinimumEventTime(minEventTime.toBuilder().setSeconds(1002).build()) - .setMinimumPublishTime(minPublishTime) - .build(); - - when(mockClient.computeMessageStats( - example(TopicPath.class), - example(Partition.class), - example(Offset.class), - Offset.of(Long.MAX_VALUE))) - .thenReturn(ApiFutures.immediateFuture(response)); - - assertEquals(reader.computeMessageStats(example(Offset.class)), response); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/UnboundedReaderImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/UnboundedReaderImplTest.java deleted file mode 100644 index c6663204ed7b..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/UnboundedReaderImplTest.java +++ /dev/null @@ -1,204 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertThrows; -import static org.junit.Assert.assertTrue; -import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; - -import com.google.cloud.pubsublite.Message; -import com.google.cloud.pubsublite.Offset; -import com.google.cloud.pubsublite.internal.testing.FakeApiService; -import com.google.cloud.pubsublite.proto.ComputeMessageStatsResponse; -import com.google.cloud.pubsublite.proto.SequencedMessage; -import com.google.protobuf.ByteString; -import com.google.protobuf.util.Timestamps; -import java.io.IOException; -import java.util.NoSuchElementException; -import java.util.Optional; -import java.util.concurrent.TimeUnit; -import org.apache.beam.sdk.io.UnboundedSource; -import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.joda.time.Instant; -import org.junit.Before; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.Mock; -import org.mockito.Spy; -import org.mockito.junit.MockitoJUnit; -import org.mockito.junit.MockitoRule; -import org.mockito.quality.Strictness; -import org.mockito.stubbing.Answer; - -@RunWith(JUnit4.class) -public class UnboundedReaderImplTest { - - private static final Offset INITIAL_OFFSET = Offset.of(1); - - @Rule public MockitoRule mockito = MockitoJUnit.rule().strictness(Strictness.STRICT_STUBS); - @Spy private UnboundedSource source; - - abstract static class FakeSubscriber extends FakeApiService implements MemoryBufferedSubscriber {} - - @Spy private FakeSubscriber subscriber; - @Mock private TopicBacklogReader backlogReader; - @Mock private BlockingCommitter committer; - - private UnboundedReaderImpl reader; - - private static SequencedMessage messageWith(String data, long offset, Instant timestamp) { - return com.google.cloud.pubsublite.SequencedMessage.of( - Message.builder().setData(ByteString.copyFromUtf8(data)).build(), - Timestamps.fromMillis(timestamp.getMillis()), - Offset.of(offset), - 21) - .toProto(); - } - - private void startSubscriber() throws Exception { - doReturn(Optional.empty()).when(subscriber).peek(); - assertFalse(reader.start()); - verify(subscriber).startAsync(); - verify(subscriber).awaitRunning(1, TimeUnit.MINUTES); - } - - private void advancePastMessage(long offset) throws Exception { - SequencedMessage message = messageWith("abc", offset, Instant.now()); - doReturn(Optional.of(message)).when(subscriber).peek(); - assertTrue(reader.advance()); - doAnswer( - (Answer) - args -> { - doReturn(Optional.empty()).when(subscriber).peek(); - return null; - }) - .when(subscriber) - .pop(); - assertFalse(reader.advance()); - } - - @Before - public void setUp() { - doReturn(INITIAL_OFFSET).when(subscriber).fetchOffset(); - reader = - new UnboundedReaderImpl(source, subscriber, backlogReader, () -> committer, Offset.of(1)); - } - - @Test - public void startAdvances() throws Exception { - Instant ts = Instant.now(); - SequencedMessage message = messageWith("abc", 2, ts); - doReturn(Optional.of(message)).when(subscriber).peek(); - assertTrue(reader.start()); - verify(subscriber).startAsync(); - verify(subscriber).awaitRunning(1, TimeUnit.MINUTES); - assertEquals(reader.getCurrent(), message); - assertEquals(reader.getWatermark(), ts); - } - - @Test - public void startAdvancesNoMessage() throws Exception { - doReturn(Optional.empty()).when(subscriber).peek(); - assertFalse(reader.start()); - verify(subscriber).startAsync(); - verify(subscriber).awaitRunning(1, TimeUnit.MINUTES); - assertThrows(NoSuchElementException.class, reader::getCurrent); - assertEquals(BoundedWindow.TIMESTAMP_MIN_VALUE, reader.getWatermark()); - } - - @Test - public void advanceNoPreviousValue() throws Exception { - startSubscriber(); - SequencedMessage message = messageWith("abc", 2, Instant.now()); - doReturn(Optional.of(message)).when(subscriber).peek(); - assertTrue(reader.advance()); - verify(subscriber, times(0)).pop(); - } - - @Test - public void advanceWithPreviousValue() throws Exception { - startSubscriber(); - Instant ts1 = Instant.now(); - Instant ts2 = Instant.now(); - SequencedMessage message1 = messageWith("abc", 2, ts1); - SequencedMessage message2 = messageWith("def", 3, ts2); - doReturn(Optional.of(message1)).when(subscriber).peek(); - assertTrue(reader.advance()); - assertEquals(reader.getCurrent(), message1); - assertEquals(reader.getCurrentTimestamp(), ts1); - assertEquals(reader.getWatermark(), ts1); - doAnswer( - (Answer) - args -> { - doReturn(Optional.of(message2)).when(subscriber).peek(); - return null; - }) - .when(subscriber) - .pop(); - assertTrue(reader.advance()); - verify(subscriber).pop(); - assertEquals(reader.getCurrent(), message2); - assertEquals(reader.getCurrentTimestamp(), ts2); - assertEquals(reader.getWatermark(), ts1); - } - - @Test - public void advanceSubscriberNotRunningThrows() throws Exception { - startSubscriber(); - subscriber.fail(new RuntimeException("I failed")); - assertThrows(IOException.class, reader::advance); - } - - @Test - public void getCheckpointMark() throws Exception { - startSubscriber(); - advancePastMessage(2); - CheckpointMarkImpl mark = reader.getCheckpointMark(); - verify(subscriber).rebuffer(); - assertEquals(3, mark.offset.value()); - } - - @Test - public void getSplitBacklogBytes() throws Exception { - startSubscriber(); - advancePastMessage(2); - doReturn(ComputeMessageStatsResponse.newBuilder().setMessageBytes(42).build()) - .when(backlogReader) - .computeMessageStats(Offset.of(3)); - assertEquals(42, reader.getSplitBacklogBytes()); - } - - @Test - public void closeClosesAll() throws Exception { - startSubscriber(); - doThrow(new IllegalStateException("abc")).when(subscriber).awaitTerminated(1, TimeUnit.MINUTES); - doThrow(new IllegalStateException("def")).when(backlogReader).close(); - assertThrows(IOException.class, reader::close); - verify(subscriber).stopAsync(); - verify(subscriber).awaitTerminated(1, TimeUnit.MINUTES); - verify(backlogReader).close(); - } -} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/UuidDeduplicationTransformTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/UuidDeduplicationTransformTest.java deleted file mode 100644 index d390cc5ea127..000000000000 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/internal/UuidDeduplicationTransformTest.java +++ /dev/null @@ -1,149 +0,0 @@ -/* - * 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.sdk.io.gcp.pubsublite.internal; - -import com.google.cloud.pubsublite.proto.AttributeValues; -import com.google.cloud.pubsublite.proto.Cursor; -import com.google.cloud.pubsublite.proto.PubSubMessage; -import com.google.cloud.pubsublite.proto.SequencedMessage; -import com.google.protobuf.ByteString; -import com.google.protobuf.util.Timestamps; -import org.apache.beam.sdk.extensions.protobuf.ProtoCoder; -import org.apache.beam.sdk.io.gcp.pubsublite.UuidDeduplicationOptions; -import org.apache.beam.sdk.testing.PAssert; -import org.apache.beam.sdk.testing.TestPipeline; -import org.apache.beam.sdk.testing.TestStream; -import org.apache.beam.sdk.transforms.Deduplicate; -import org.apache.beam.sdk.values.PCollection; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; -import org.joda.time.Duration; -import org.joda.time.Instant; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class UuidDeduplicationTransformTest { - @Rule public final TestPipeline pipeline = TestPipeline.create(); - private static final Instant START = new Instant(0); - - private static SequencedMessage newMessage() { - Uuid uuid = Uuid.random(); - return SequencedMessage.newBuilder() - .setMessage( - PubSubMessage.newBuilder() - .putAttributes( - Uuid.DEFAULT_ATTRIBUTE, - AttributeValues.newBuilder().addValues(uuid.value()).build())) - .setSizeBytes(10000) - .setPublishTime(Timestamps.EPOCH) - .setCursor(Cursor.newBuilder().setOffset(10)) - .build(); - } - - @Test - public void unrelatedUuidsProxied() { - SequencedMessage message1 = newMessage(); - SequencedMessage message2 = newMessage(); - - TestStream messageStream = - TestStream.create(ProtoCoder.of(SequencedMessage.class)) - .advanceWatermarkTo(START) - .addElements(message1) - .advanceWatermarkTo(START.plus(Deduplicate.DEFAULT_DURATION.dividedBy(2))) - .addElements(message2) - .advanceWatermarkToInfinity(); - PCollection results = - pipeline - .apply(messageStream) - .apply(new UuidDeduplicationTransform(UuidDeduplicationOptions.newBuilder().build())); - PAssert.that(results).containsInAnyOrder(message1, message2); - pipeline.run(); - } - - @Test - public void sameUuidsWithinWindowOnlyOne() { - SequencedMessage message = newMessage(); - - TestStream messageStream = - TestStream.create(ProtoCoder.of(SequencedMessage.class)) - .advanceWatermarkTo(START) - .addElements(message) - .advanceWatermarkTo(START.plus(Deduplicate.DEFAULT_DURATION.dividedBy(2))) - .advanceWatermarkToInfinity(); - PCollection results = - pipeline - .apply(messageStream) - .apply(new UuidDeduplicationTransform(UuidDeduplicationOptions.newBuilder().build())); - PAssert.that(results).containsInAnyOrder(message); - pipeline.run(); - } - - @Test - public void sameUuidsAfterGcOutsideWindowHasBoth() { - SequencedMessage message1 = newMessage(); - - TestStream messageStream = - TestStream.create(ProtoCoder.of(SequencedMessage.class)) - .advanceWatermarkTo(START) - .addElements(message1) - .advanceWatermarkTo( - START.plus( - UuidDeduplicationOptions.DEFAULT_DEDUPLICATE_DURATION.plus(Duration.millis(1)))) - .addElements(message1) - .advanceWatermarkToInfinity(); - PCollection results = - pipeline - .apply(messageStream) - .apply(new UuidDeduplicationTransform(UuidDeduplicationOptions.newBuilder().build())); - PAssert.that(results).containsInAnyOrder(message1, message1); - pipeline.run(); - } - - @Test - public void dedupesBasedOnReturnedUuid() { - byte[] bytes = {(byte) 0x123, (byte) 0x456}; - // These messages have different uuids, so they would both appear in the output collection if - // the extractor is not respected. - SequencedMessage message1 = newMessage(); - SequencedMessage message2 = newMessage(); - - TestStream messageStream = - TestStream.create(ProtoCoder.of(SequencedMessage.class)) - .advanceWatermarkTo(START) - .addElements(message1, message2) - .advanceWatermarkToInfinity(); - PCollection results = - pipeline - .apply(messageStream) - .apply( - new UuidDeduplicationTransform( - UuidDeduplicationOptions.newBuilder() - .setUuidExtractor(message -> Uuid.of(ByteString.copyFrom(bytes))) - .build())); - PAssert.that(results) - .satisfies( - messages -> { - Preconditions.checkArgument(Iterables.size(messages) == 1); - return null; - }); - pipeline.run(); - } -} diff --git a/sdks/python/apache_beam/io/gcp/pubsublite/__init__.py b/sdks/python/apache_beam/io/gcp/pubsublite/__init__.py deleted file mode 100644 index 565777e14050..000000000000 --- a/sdks/python/apache_beam/io/gcp/pubsublite/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -# -# 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. -# - -from .proto_api import ReadFromPubSubLite -from .proto_api import WriteToPubSubLite - -__all__ = [ - "ReadFromPubSubLite", - "WriteToPubSubLite", -] diff --git a/sdks/python/apache_beam/io/gcp/pubsublite/external.py b/sdks/python/apache_beam/io/gcp/pubsublite/external.py deleted file mode 100644 index a0e46c1b4d88..000000000000 --- a/sdks/python/apache_beam/io/gcp/pubsublite/external.py +++ /dev/null @@ -1,115 +0,0 @@ -# -# 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. -# - -"""Google Pub/Sub Lite sources and sinks. - -This API is currently under development and is subject to change. -""" - -# pytype: skip-file - -import typing - -from apache_beam.transforms.external import BeamJarExpansionService -from apache_beam.transforms.external import ExternalTransform -from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder - -_ReadSchema = typing.NamedTuple( - '_ReadSchema', [('subscription_path', str), ('deduplicate', bool)]) - - -def _default_io_expansion_service(): - return BeamJarExpansionService( - 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar') - - -class _ReadExternal(ExternalTransform): - """ - An external PTransform which reads from Pub/Sub Lite and returns a - SequencedMessage as serialized bytes. - - This transform is not part of the public API. - - Experimental; no backwards-compatibility guarantees. - """ - def __init__( - self, - subscription_path, - deduplicate=None, - expansion_service=None, - ): - """ - Initializes a read operation from Pub/Sub Lite, returning the serialized - bytes of SequencedMessage protos. - - Args: - subscription_path: A Pub/Sub Lite Subscription path. - deduplicate: Whether to deduplicate messages based on the value of - the 'x-goog-pubsublite-dataflow-uuid' attribute. - """ - if deduplicate is None: - deduplicate = False - if expansion_service is None: - expansion_service = _default_io_expansion_service() - super().__init__( - 'beam:transform:org.apache.beam:pubsublite_read:v1', - NamedTupleBasedPayloadBuilder( - _ReadSchema( - subscription_path=subscription_path, deduplicate=deduplicate)), - expansion_service) - - -_WriteSchema = typing.NamedTuple( - '_WriteSchema', [('topic_path', str), ('add_uuids', bool)]) - - -class _WriteExternal(ExternalTransform): - """ - An external PTransform which writes serialized PubSubMessage protos to - Pub/Sub Lite. - - This transform is not part of the public API. - - Experimental; no backwards-compatibility guarantees. - """ - def __init__( - self, - topic_path, - add_uuids=None, - expansion_service=None, - ): - """ - Initializes a write operation to Pub/Sub Lite, writing the serialized bytes - of PubSubMessage protos. - - Args: - topic_path: A Pub/Sub Lite Topic path. - add_uuids: Whether to add uuids to the 'x-goog-pubsublite-dataflow-uuid' - uuid attribute. - """ - if add_uuids is None: - add_uuids = False - if expansion_service is None: - expansion_service = _default_io_expansion_service() - super().__init__( - 'beam:transform:org.apache.beam:pubsublite_write:v1', - NamedTupleBasedPayloadBuilder( - _WriteSchema( - topic_path=topic_path, - add_uuids=add_uuids, - )), - expansion_service) diff --git a/sdks/python/apache_beam/io/gcp/pubsublite/proto_api.py b/sdks/python/apache_beam/io/gcp/pubsublite/proto_api.py deleted file mode 100644 index a8e3defc4f99..000000000000 --- a/sdks/python/apache_beam/io/gcp/pubsublite/proto_api.py +++ /dev/null @@ -1,106 +0,0 @@ -# -# 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. -# - -from apache_beam.io.gcp.pubsublite.external import _ReadExternal -from apache_beam.io.gcp.pubsublite.external import _WriteExternal -from apache_beam.transforms import Map -from apache_beam.transforms import PTransform - -try: - from google.cloud import pubsublite -except ImportError: - pubsublite = None - - -class ReadFromPubSubLite(PTransform): - """ - A ``PTransform`` for reading from Pub/Sub Lite. - - Produces a PCollection of google.cloud.pubsublite.SequencedMessage - - Experimental; no backwards-compatibility guarantees. - """ - def __init__( - self, - subscription_path, - deduplicate=None, - expansion_service=None, - ): - """Initializes ``ReadFromPubSubLite``. - - Args: - subscription_path: Pub/Sub Lite Subscription in the form - projects//locations//subscriptions/ - deduplicate: Whether to deduplicate messages based on the value of - the 'x-goog-pubsublite-dataflow-uuid' attribute. Defaults to False. - """ - super().__init__() - self._source = _ReadExternal( - subscription_path=subscription_path, - deduplicate=deduplicate, - expansion_service=expansion_service, - ) - - def expand(self, pvalue): - pcoll = pvalue.pipeline | self._source - pcoll.element_type = bytes - pcoll = pcoll | Map(pubsublite.SequencedMessage.deserialize) - pcoll.element_type = pubsublite.SequencedMessage - return pcoll - - -class WriteToPubSubLite(PTransform): - """ - A ``PTransform`` for writing to Pub/Sub Lite. - - Consumes a PCollection of google.cloud.pubsublite.PubSubMessage - - Experimental; no backwards-compatibility guarantees. - """ - def __init__( - self, - topic_path, - add_uuids=None, - expansion_service=None, - ): - """Initializes ``WriteToPubSubLite``. - - Args: - topic_path: A Pub/Sub Lite Topic path. - add_uuids: Whether to add uuids to the 'x-goog-pubsublite-dataflow-uuid' - uuid attribute. Defaults to False. - """ - super().__init__() - self._source = _WriteExternal( - topic_path=topic_path, - add_uuids=add_uuids, - expansion_service=expansion_service, - ) - - @staticmethod - def _message_to_proto_str(element: pubsublite.PubSubMessage): - if not isinstance(element, pubsublite.PubSubMessage): - raise TypeError( - 'Unexpected element. Type: %s (expected: PubSubMessage), ' - 'value: %r' % (type(element), element)) - return pubsublite.PubSubMessage.serialize(element) - - def expand(self, pcoll): - pcoll = pcoll | Map(WriteToPubSubLite._message_to_proto_str) - pcoll.element_type = bytes - pcoll = pcoll | self._source - return pcoll diff --git a/sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/src/yaml/EmojiMap.ts b/sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/src/yaml/EmojiMap.ts index ed6a9f2285c8..800ec6aa9eea 100644 --- a/sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/src/yaml/EmojiMap.ts +++ b/sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/src/yaml/EmojiMap.ts @@ -62,8 +62,6 @@ export const transformEmojiMap: Record = { WriteToPostgres: '⬆️🐘', ReadFromPubSub: '⬇️📢', WriteToPubSub: '⬆️📢', - ReadFromPubSubLite: '⬇️📣', - WriteToPubSubLite: '⬆️📣', ReadFromSpanner: '⬇️📏', WriteToSpanner: '⬆️📏', ReadFromSqlServer: '⬇️🗄️', diff --git a/sdks/python/container/ml/py310/base_image_requirements.txt b/sdks/python/container/ml/py310/base_image_requirements.txt index 729ab317e648..673e63f0ac95 100644 --- a/sdks/python/container/ml/py310/base_image_requirements.txt +++ b/sdks/python/container/ml/py310/base_image_requirements.txt @@ -82,7 +82,6 @@ google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-profiler==4.1.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/container/ml/py310/gpu_image_requirements.txt b/sdks/python/container/ml/py310/gpu_image_requirements.txt index ad5095bec5eb..553de91f9ef3 100644 --- a/sdks/python/container/ml/py310/gpu_image_requirements.txt +++ b/sdks/python/container/ml/py310/gpu_image_requirements.txt @@ -99,7 +99,6 @@ google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-profiler==4.1.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/container/ml/py311/base_image_requirements.txt b/sdks/python/container/ml/py311/base_image_requirements.txt index f078095f6948..cfdf71ce77bd 100644 --- a/sdks/python/container/ml/py311/base_image_requirements.txt +++ b/sdks/python/container/ml/py311/base_image_requirements.txt @@ -80,7 +80,6 @@ google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-profiler==4.1.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/container/ml/py311/gpu_image_requirements.txt b/sdks/python/container/ml/py311/gpu_image_requirements.txt index 1fd5d768406f..313cad9147bd 100644 --- a/sdks/python/container/ml/py311/gpu_image_requirements.txt +++ b/sdks/python/container/ml/py311/gpu_image_requirements.txt @@ -97,7 +97,6 @@ google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-profiler==4.1.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/container/ml/py312/base_image_requirements.txt b/sdks/python/container/ml/py312/base_image_requirements.txt index 0bbe666b6805..da889c0526d0 100644 --- a/sdks/python/container/ml/py312/base_image_requirements.txt +++ b/sdks/python/container/ml/py312/base_image_requirements.txt @@ -79,7 +79,6 @@ google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-profiler==4.1.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/container/ml/py312/gpu_image_requirements.txt b/sdks/python/container/ml/py312/gpu_image_requirements.txt index d0024977e5c0..ca40f87da0b2 100644 --- a/sdks/python/container/ml/py312/gpu_image_requirements.txt +++ b/sdks/python/container/ml/py312/gpu_image_requirements.txt @@ -96,7 +96,6 @@ google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-profiler==4.1.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/container/ml/py313/base_image_requirements.txt b/sdks/python/container/ml/py313/base_image_requirements.txt index 1a2e88342d24..5e3b32e82bd4 100644 --- a/sdks/python/container/ml/py313/base_image_requirements.txt +++ b/sdks/python/container/ml/py313/base_image_requirements.txt @@ -77,7 +77,6 @@ google-cloud-kms==3.9.0 google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/container/py310/base_image_requirements.txt b/sdks/python/container/py310/base_image_requirements.txt index cc61bef566fe..61c827e09630 100644 --- a/sdks/python/container/py310/base_image_requirements.txt +++ b/sdks/python/container/py310/base_image_requirements.txt @@ -76,7 +76,6 @@ google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-profiler==4.1.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/container/py311/base_image_requirements.txt b/sdks/python/container/py311/base_image_requirements.txt index 09e6e9dc1453..7be8ec580698 100644 --- a/sdks/python/container/py311/base_image_requirements.txt +++ b/sdks/python/container/py311/base_image_requirements.txt @@ -74,7 +74,6 @@ google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-profiler==4.1.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/container/py312/base_image_requirements.txt b/sdks/python/container/py312/base_image_requirements.txt index 398a4b8285c7..f2c9ce2786f2 100644 --- a/sdks/python/container/py312/base_image_requirements.txt +++ b/sdks/python/container/py312/base_image_requirements.txt @@ -73,7 +73,6 @@ google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-profiler==4.1.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/container/py313/base_image_requirements.txt b/sdks/python/container/py313/base_image_requirements.txt index d5b9331a9be2..6a919da4bc93 100644 --- a/sdks/python/container/py313/base_image_requirements.txt +++ b/sdks/python/container/py313/base_image_requirements.txt @@ -71,7 +71,6 @@ google-cloud-kms==3.9.0 google-cloud-language==2.19.0 google-cloud-monitoring==2.29.0 google-cloud-pubsub==2.34.0 -google-cloud-pubsublite==1.13.0 google-cloud-recommendations-ai==0.10.18 google-cloud-resource-manager==1.16.0 google-cloud-secret-manager==2.26.0 diff --git a/sdks/python/setup.py b/sdks/python/setup.py index bcdd1d77584a..0e27d99deb3b 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -479,7 +479,6 @@ def get_portability_package_data(): 'google-auth-httplib2>=0.1.0,<0.3.0', 'google-cloud-datastore>=2.0.0,<3', 'google-cloud-pubsub>=2.1.0,<3', - 'google-cloud-pubsublite>=1.2.0,<2', 'google-cloud-storage>=2.18.2,<3', # GCP packages required by tests 'google-cloud-bigquery>=2.0.0,<4',