Examples with ElementByteSizeObserver org.apache.beam.sdk.util.common.ElementByteSizeObserver used on opensource projects

Example 1 with ElementByteSizeObserver

use of org.apache.beam.sdk.util.common.ElementByteSizeObserver in project beam by apache.

the class PartialGroupByKeyParDoFnsTest method testCoderSizeEstimationWithLazyObserver.

@Test
public void testCoderSizeEstimationWithLazyObserver() throws Exception {
    // Have the code report that byte size observation is lazy
    doAnswer(invocation -> {
        Object[] args = invocation.getArguments();
        // Set lazy to true and do not update the size at all
        ((ElementByteSizeObserver) args[1]).setLazy();
        return null;
    }).when(mockCoder).registerByteSizeObserver(Matchers.eq("apple"), Matchers.<ElementByteSizeObserver>any());
    // Encode the input to the output stream
    doAnswer(invocation -> {
        Object[] args = invocation.getArguments();
        String value = (String) args[0];
        OutputStream os = (OutputStream) args[1];
        os.write(value.getBytes(StandardCharsets.UTF_8));
        return null;
    }).when(mockCoder).encode(Matchers.eq("apple"), Matchers.<OutputStream>any());
    CoderSizeEstimator<String> estimator = new CoderSizeEstimator(mockCoder);
    // Observer never updates size, so if result is 5, must have delegated to actual encoding
    assertEquals(5L, estimator.estimateSize("apple"));
}