onBodySent is not called in mockDispatch

[tjhiggins]

Bug Description

onBodySent is not called when using the MockAgent

Reproducible By

import { DecoratorHandler, Dispatcher, fetch, MockAgent } from 'undici';

class LoggingHandler extends DecoratorHandler implements Dispatcher.DispatchHandlers {
  private requestBody = '';

  constructor(
    protected requestOptions: Dispatcher.DispatchOptions,
    protected handler: Dispatcher.DispatchHandlers,
  ) {
    super(handler);
  }

  onConnect(abort: (err?: Error) => void): void {
    if (this.handler.onConnect) {
      return this.handler.onConnect(abort);
    }
  }

  onBodySent(chunk: any, totalBytesSent: number) {
    try {
      this.requestBody += Buffer.from(chunk).toString();
    } catch (err) {
      console.debug('Failed to parse request body:', err);
    }
    if (this.handler.onBodySent) {
      this.handler.onBodySent(chunk, totalBytesSent);
    }
  }

  onComplete(trailers: string[] | null) {
    console.log('onComplete', this.requestBody); // In reality this would log to something like GCP with headers/response/status etc
    return this.handler.onComplete!(trailers);
  }
}

function createLoggingInterceptor() {
  return (dispatch: Dispatcher['dispatch']) => {
    return function Cache(requestOptions: Dispatcher.DispatchOptions, handler: Dispatcher.DispatchHandlers) {
      return dispatch(requestOptions, new LoggingHandler(requestOptions, handler));
    };
  };
}

describe('onBodySent', () => {
  it('onBodySent should be called when using MockAgent', async () => {
    const mockAgent = new MockAgent();
    mockAgent.disableNetConnect();
    const mockPool = mockAgent.get('http://mock.localhost');

    const onConnectSpy = jest.spyOn(LoggingHandler.prototype, 'onConnect');
    const onBodySentSpy = jest.spyOn(LoggingHandler.prototype, 'onBodySent');

    const dispatcher = mockAgent.compose(createLoggingInterceptor());

    const bodyString = JSON.stringify({
      amount: '100',
    });

    mockPool
      .intercept({
        path: '/test',
        method: 'POST',
        body: () => {
          return true;
        },
      })
      .reply(200, {
        test: 1,
      });

    const res = await fetch('http://mock.localhost/test', {
      dispatcher,
      method: 'POST',
      body: bodyString,
    });
    expect(await res.json()).toEqual({ test: 1 });
    expect(onConnectSpy).toHaveBeenCalledTimes(1);
    expect(onBodySentSpy).toHaveBeenCalledTimes(1); // fails
  });
});

Expected Behavior

onBodySent would be called when making a POST request

Additional context

Looks like it should be called here:

undici/lib/mock/mock-utils.js

Line 317 in a73fd09

handler.onHeaders?.(statusCode, responseHeaders, resume, getStatusText(statusCode))

onBodySent future?

Concerned about these future changes which remove onBodySent. Not sure how I'd easily get the request body in an interceptor in the next branch. Mostly creating this issue to figure out what I should be doing instead going forward.
#2723 (comment)

[mcollina]

Thanks for reporting! Would you like to send a Pull Request to address this issue? Remember to add unit tests.

[metcoder95]

FWIW, This can be achieved in the same way with the current interceptor setup. e.g.

return (dispatch: Dispatcher['dispatch']) => {
    return function intercept(requestOptions: Dispatcher.DispatchOptions, handler: Dispatcher.DispatchHandlers) {
	  // if body is stream, attach to it
      requestOptions.body.on('data', () => { /* do something */ })

      return dispatch(requestOptions, new LoggingHandler(requestOptions, handler));
    };

[tjhiggins]

FWIW, This can be achieved in the same way with the current interceptor setup. e.g.

return (dispatch: Dispatcher['dispatch']) => {
    return function intercept(requestOptions: Dispatcher.DispatchOptions, handler: Dispatcher.DispatchHandlers) {
	  // if body is stream, attach to it
      requestOptions.body.on('data', () => { /* do something */ })

      return dispatch(requestOptions, new LoggingHandler(requestOptions, handler));
    };

@metcoder95 I appreciate the quick response. I don't think this works for fetch since it's requestOptions.body is an AsyncGenerator:

undici/lib/web/fetch/index.js

Line 1845 in 444a0af

requestBody = (async function * () {

I tried wrapping in a Readable.from, but ran into other issues with content-length being set in the headers.

[tjhiggins]

This feels pretty bad, but it seems to work:

export function createLoggingInterceptor({ projectId, name }: LoggingOptions) {
  const logging = new Logging({ projectId, autoRetry: true, maxRetries: 3 });
  const log = logging.log(name);
  return (dispatch: Dispatcher['dispatch']) => {
    return function Logging(requestOptions: Dispatcher.DispatchOptions, handler: Dispatcher.DispatchHandlers) {
      let body1: PassThrough | undefined;
      let body2: PassThrough | undefined;
      if (requestOptions.body) {
        body1 = new PassThrough();
        body2 = new PassThrough();
        const body = Readable.from(requestOptions.body); // requestOptions.body is an AsyncGenerator and can only be read once
        body.pipe(body1);
        body.pipe(body2);
      }
      return dispatch(
        { ...requestOptions, body: body1 },
        new LoggingHandler({ log }, { ...requestOptions, body: body2 }, handler),
      );
    };
  };
}

[metcoder95]

Yeah, AsyncGenerators cannot be re-used/consumed in parallel.

Can you elaborate why bad? The interceptors API has been made for these kind of purposes