public static <InputT, OutputT> DoFnInvoker<InputT, OutputT> invokerFor(
      DoFn<InputT, OutputT> fn) {
    return ByteBuddyDoFnInvokerFactory.only().newByteBuddyInvoker(fn);
  }
@Test
  public void testSplittableDoFnWithAllMethods() throws Exception {
    MockFn fn = mock(MockFn.class);
    DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
    final SomeRestrictionTracker tracker = mock(SomeRestrictionTracker.class);
    final SomeRestrictionCoder coder = mock(SomeRestrictionCoder.class);
    final InstantCoder watermarkEstimatorStateCoder = InstantCoder.of();
    final Instant watermarkEstimatorState = Instant.now();
    final WatermarkEstimator<Instant> watermarkEstimator =
        new WatermarkEstimators.Manual(watermarkEstimatorState);
    SomeRestriction restriction = new SomeRestriction();
    final SomeRestriction part1 = new SomeRestriction();
    final SomeRestriction part2 = new SomeRestriction();
    final SomeRestriction part3 = new SomeRestriction();
    when(fn.getRestrictionCoder()).thenReturn(coder);
    when(fn.getWatermarkEstimatorStateCoder()).thenReturn(watermarkEstimatorStateCoder);
    when(fn.getInitialRestriction(mockElement)).thenReturn(restriction);
    doAnswer(
            AdditionalAnswers.delegatesTo(
                new MockFn() {
                  @DoFn.SplitRestriction
                  @Override
                  public void splitRestriction(
                      @Element String element,
                      @Restriction SomeRestriction restriction,
                      DoFn.OutputReceiver<SomeRestriction> receiver) {
                    receiver.output(part1);
                    receiver.output(part2);
                    receiver.output(part3);
                  }
                }))
        .when(fn)
        .splitRestriction(eq(mockElement), same(restriction), any());
    when(fn.getInitialWatermarkEstimatorState()).thenReturn(watermarkEstimatorState);
    when(fn.newTracker(restriction)).thenReturn(tracker);
    when(fn.newWatermarkEstimator(watermarkEstimatorState)).thenReturn(watermarkEstimator);
    when(fn.processElement(mockProcessContext, tracker, watermarkEstimator)).thenReturn(resume());
    when(fn.getSize()).thenReturn(2.0);

    assertEquals(coder, invoker.invokeGetRestrictionCoder(CoderRegistry.createDefault()));
    assertEquals(
        watermarkEstimatorStateCoder,
        invoker.invokeGetWatermarkEstimatorStateCoder(CoderRegistry.createDefault()));
    assertEquals(
        restriction,
        invoker.invokeGetInitialRestriction(
            new FakeArgumentProvider<String, String>() {
              @Override
              public String element(DoFn<String, String> doFn) {
                return mockElement;
              }
            }));
    List<SomeRestriction> outputs = new ArrayList<>();
    invoker.invokeSplitRestriction(
        new FakeArgumentProvider<String, String>() {
          @Override
          public String element(DoFn<String, String> doFn) {
            return mockElement;
          }

          @Override
          public Object restriction() {
            return restriction;
          }

          @Override
          public OutputReceiver outputReceiver(DoFn doFn) {
            return new OutputReceiver<SomeRestriction>() {
              @Override
              public void output(SomeRestriction output) {
                outputs.add(output);
              }

              @Override
              public void outputWithTimestamp(SomeRestriction output, Instant timestamp) {
                fail("Unexpected output with timestamp");
              }

              @Override
              public void outputWindowedValue(
                  SomeRestriction output,
                  Instant timestamp,
                  Collection<? extends BoundedWindow> windows,
                  PaneInfo paneInfo) {
                fail("Unexpected outputWindowedValue");
              }
            };
          }
        });

    assertEquals(Arrays.asList(part1, part2, part3), outputs);
    assertEquals(
        watermarkEstimatorState,
        invoker.invokeGetInitialWatermarkEstimatorState(new FakeArgumentProvider<>()));
    assertEquals(
        tracker,
        invoker.invokeNewTracker(
            new FakeArgumentProvider<String, String>() {
              @Override
              public String element(DoFn<String, String> doFn) {
                return mockElement;
              }

              @Override
              public Object restriction() {
                return restriction;
              }
            }));
    assertEquals(
        watermarkEstimator,
        invoker.invokeNewWatermarkEstimator(
            new FakeArgumentProvider<String, String>() {
              @Override
              public Object watermarkEstimatorState() {
                return watermarkEstimatorState;
              }
            }));
    assertEquals(
        resume(),
        invoker.invokeProcessElement(
            new FakeArgumentProvider<String, String>() {
              @Override
              public DoFn<String, String>.ProcessContext processContext(DoFn<String, String> fn) {
                return mockProcessContext;
              }

              @Override
              public RestrictionTracker<?, ?> restrictionTracker() {
                return tracker;
              }

              @Override
              public WatermarkEstimator<?> watermarkEstimator() {
                return watermarkEstimator;
              }
            }));
    assertEquals(2.0, invoker.invokeGetSize(mockArgumentProvider), 0.0001);
  }