Mocking InvokeModelWithResponseStreamOutput with eventStream

[mousedownmike]

I'm trying to create unit tests where I call InvokeModelWithResposeStream. To do this I need to have a InvokeModelWithResponseStreamOutput that contains a InvokeModelWithResponseStreamEventStream that I can mock. There is a helper that I am able to use to create that EventStream but there is no way to actually return it from the StreamOutput (eventStream is not exported).

The SDK seems to indicate that this should be possible with the following comment, however, I'm unable to construct an output that returns the mocked stream. I'm not sure what I'm missing:

// InvokeModelWithResponseStreamEventStream provides the event stream handling for the InvokeModelWithResponseStream operation.
//
// For testing and mocking the event stream this type should be initialized via
// the NewInvokeModelWithResponseStreamEventStream constructor function. Using the functional options
// to pass in nested mock behavior.

[mousedownmike]

I ended up creating an adapter that allowed me to return the mocked EventStream. It's not ideal, but it gets the job done.

type bedrockClient interface {
	InvokeModelWithResponseStream(ctx context.Context, params *bedrockruntime.InvokeModelWithResponseStreamInput, optFns ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelWithResponseStreamOutput, error)
}

type modelInvoker interface {
	InvokeModelWithResponseStream(ctx context.Context, params *bedrockruntime.InvokeModelWithResponseStreamInput, optFns ...func(*bedrockruntime.Options)) (types.InvokeModelWithResponseStreamOutput, error)
}

// The bedrockClientAdapter allows us to mock the EventStream output
type bedrockClientAdapter struct {
	client bedrockClient
}

func newBedrockClientAdapter(client bedrockClient) *bedrockClientAdapter {
	return &bedrockClientAdapter{client: client}
}

func (a *bedrockClientAdapter) InvokeModelWithResponseStream(ctx context.Context, params *bedrockruntime.InvokeModelWithResponseStreamInput, optFns ...func(*bedrockruntime.Options)) (types.InvokeModelWithResponseStreamOutput, error) {
	return a.client.InvokeModelWithResponseStream(ctx, params, optFns...)
}