1import { Message } from "google-protobuf";
2import { grpc } from "@improbable-eng/grpc-web";
3import assignIn = require("lodash.assignin");
4
5function frameResponse(request: Message): Uint8Array {
6  const bytes = request.serializeBinary();
7  const frame = new ArrayBuffer(bytes.byteLength + 5);
8  new DataView(frame, 0, 5).setUint32(1, bytes.length, false /* big endian */);
9  new Uint8Array(frame, 5).set(bytes);
10  return new Uint8Array(frame);
11}
12
13export function frameRequest(request: Message): ArrayBufferView {
14  const bytes = request.serializeBinary();
15  const frame = new ArrayBuffer(bytes.byteLength + 5);
16  new DataView(frame, 1, 4).setUint32(0, bytes.length, false /* big endian */);
17  new Uint8Array(frame, 5).set(bytes);
18  return new Uint8Array(frame);
19}
20
21function frameTrailers(trailers: grpc.Metadata): Uint8Array {
22  let asString = "";
23  trailers.forEach((key: string, values: string[]) => {
24    asString += `${key}: ${values.join(", ")}\r\n`;
25  });
26  const bytes = new Buffer(asString);
27  const frame = new ArrayBuffer(bytes.byteLength + 5);
28  const dataview = new DataView(frame, 0, 5);
29  dataview.setUint32(1, bytes.length, false /* big endian */);
30  dataview.setUint8(0, 128);
31  new Uint8Array(frame, 5).set(bytes);
32  return new Uint8Array(frame);
33}
34
35/**
36 * FakeTransportBuilder is a builder pattern implementation which allows you to configure a new
37 * FakeTransport instance.
38 *
39 * @class
40 */
41export class FakeTransportBuilder {
42  private requestListener: (options: grpc.TransportOptions) => void;
43  private headersListener: (headers: grpc.Metadata) => void;
44  private messageListener: (messageBytes: ArrayBufferView) => void;
45  private finishSendListener: () => void;
46  private cancelListener: () => void;
47  private preHeadersErrorCode: number | null = null;
48  private headers: grpc.Metadata | null = null;
49  private preMessagesError: [grpc.Code, string] | null = null;
50  private messages: Array<Message> = [];
51  private preTrailersError: [grpc.Code, string] | null = null;
52  private trailers: grpc.Metadata | null = null;
53  private autoTrigger: boolean = true;
54
55  /**
56   * withRequestListener is a hook which allows you to listen for when a request is made to the server
57   * via the FakeTransport being constructed.
58   *
59   * @param {(options: TransportOptions) => void} requestListener
60   * @returns {this}
61   */
62  withRequestListener(requestListener: (options: grpc.TransportOptions) => void) {
63    this.requestListener = requestListener;
64    return this;
65  }
66
67  /**
68   * withHeadersListener is a hook which allows you to listen for any headers which were sent to the
69   * server via the FakeTransport being constructed.
70   *
71   * @param {(headers: BrowserHeaders) => void} headersListener
72   * @returns {this}
73   */
74  withHeadersListener(headersListener: (headers: grpc.Metadata) => void) {
75    this.headersListener = headersListener;
76    return this;
77  }
78
79  /**
80   * withMessageListener is a hook which allows you to listen for any messages sent by the client to
81   * the server via the FakeTransport being constructed.
82   *
83   * @param {(messageBytes: ArrayBufferView) => void} messageListener
84   * @returns {this}
85   */
86  withMessageListener(messageListener: (messageBytes: ArrayBufferView) => void) {
87    this.messageListener = messageListener;
88    return this;
89  }
90
91  /**
92   * withFinishSendListener is a hook which allows you to listen for when the client finishes
93   * sending messages to the server.
94   *
95   * @param {() => void} finishSendListener
96   * @returns {this}
97   */
98  withFinishSendListener(finishSendListener: () => void) {
99    this.finishSendListener = finishSendListener;
100    return this;
101  }
102
103  /**
104   * withCancelListener is a hook which allows you to listen for a request to close/cancel the
105   * FakeTransport being constructed.
106   *
107   * @param {() => void} cancelListener
108   * @returns {this}
109   */
110  withCancelListener(cancelListener: () => void) {
111    this.cancelListener = cancelListener;
112    return this;
113  }
114
115  /**
116   * withPreHeadersError allows you to simulate a fault with the server where the request was
117   * terminated before any grpc Headers could be sent.
118   *
119   * @param {number} httpStatusCode
120   * @returns {this}
121   */
122  withPreHeadersError(httpStatusCode: number) {
123    this.preHeadersErrorCode = httpStatusCode;
124    return this;
125  }
126
127  /**
128   * withPreMessagesError allows you to simulate a fault with the server were the request was terminated
129   * after the grpc Headers were sent, but before any messages could be sent.
130   *
131   * @param {Code} grpcStatus
132   * @param {string} grpcMessage
133   * @returns {this}
134   */
135  withPreMessagesError(grpcStatus: grpc.Code, grpcMessage: string) {
136    this.preMessagesError = [grpcStatus, grpcMessage];
137    return this;
138  }
139
140  /**
141   * withPreTrailersError allows you to simulate a fault with the server where the request was
142   * terminated after the grpc Headers, and messages were sent, but before the grpc Trailers could be
143   * sent.
144   *
145   * @param {Code} grpcStatus
146   * @param {string} grpcMessage
147   * @returns {this}
148   */
149  withPreTrailersError(grpcStatus: grpc.Code, grpcMessage: string) {
150    this.preTrailersError = [grpcStatus, grpcMessage];
151    return this;
152  }
153
154  /**
155   * withHeaders allows you to stub the grpc Headers which will be returned by the server in response
156   * to any request made via the FakeTransport being constructed.
157   *
158   * @param {Metadata} headers
159   * @returns {this}
160   */
161  withHeaders(headers: grpc.Metadata) {
162    this.headers = headers;
163    return this;
164  }
165
166  /**
167   * withMessages allows you to stub the messages which will be returned by the server in response
168   * to any request made via the FakeTransport being constructed.
169   *
170   * @param {Array<Message>} messages
171   * @returns {this}
172   */
173  withMessages(messages: Array<Message>) {
174    this.messages = messages;
175    return this;
176  }
177
178  /**
179   * withTrailers allows you to sub the grpc Trailers which will be returned by the server in resposne
180   * to any request made via the FakeTransport being constructed.
181   *
182   * @param {Metadata} trailers
183   * @returns {this}
184   */
185  withTrailers(trailers: grpc.Metadata) {
186    this.trailers = trailers;
187    return this;
188  }
189
190  /**
191   * withManualTrigger allows you to have control over when the headers, messages and trailers are
192   * returned from the server to the client.
193   *
194   * @returns {this}
195   */
196  withManualTrigger() {
197    this.autoTrigger = false;
198    return this;
199  }
200
201  /**
202   * build constructs and returns the FakeTransport instance.
203   * @returns {TriggerableTransport}
204   */
205  build(): TriggerableTransport {
206    const mock = this;
207
208    const triggers = {
209      options: null as (grpc.TransportOptions | null),
210      sendHeaders: () => {
211        if (!triggers.options) {
212          throw new Error("sendHeaders called before transport had been invoked");
213        }
214        if (mock.preHeadersErrorCode !== null) {
215          triggers.options.onHeaders(new grpc.Metadata(), mock.preHeadersErrorCode);
216          triggers.options.onEnd();
217          return false;
218        }
219        const headers = mock.headers || new grpc.Metadata();
220        triggers.options.onHeaders(headers, 200);
221        return true;
222      },
223      sendMessages: () => {
224        if (!triggers.options) {
225          throw new Error("sendMessages called before transport had been invoked");
226        }
227        if (mock.preMessagesError !== null) {
228          triggers.options.onHeaders(new grpc.Metadata({ "grpc-status": String(mock.preMessagesError[0]), "grpc-message": mock.preMessagesError[1] }), 200);
229          triggers.options.onEnd();
230          return false;
231        }
232
233        mock.messages.forEach(message => {
234          triggers.options!.onChunk(frameResponse(message));
235        });
236        return true;
237      },
238      sendTrailers: () => {
239        if (!triggers.options) {
240          throw new Error("sendTrailers called before transport had been invoked");
241        }
242        if (mock.preTrailersError !== null) {
243          triggers.options.onChunk(frameTrailers(new grpc.Metadata({ "grpc-status": String(mock.preTrailersError[0]), "grpc-message": mock.preTrailersError[1] })));
244          triggers.options.onEnd();
245          return false;
246        }
247
248        const trailers = mock.trailers ? mock.trailers : new grpc.Metadata();
249
250        // Explicit status OK
251        trailers.set("grpc-status", "0");
252        triggers.options.onChunk(frameTrailers(trailers));
253        triggers.options.onEnd();
254        return true;
255      },
256      sendAll: () => {
257        if (!triggers.options) {
258          throw new Error("sendAll called before transport had been invoked");
259        }
260        if (triggers.sendHeaders()) {
261          if (triggers.sendMessages()) {
262            triggers.sendTrailers();
263          }
264        }
265      },
266    };
267
268    const transportConstructor = (optionsArg: grpc.TransportOptions) => {
269      triggers.options = optionsArg;
270
271      if (mock.requestListener) {
272        mock.requestListener(optionsArg);
273      }
274
275      return {
276        start: (metadata: grpc.Metadata) => {
277          if (mock.headersListener) {
278            mock.headersListener(metadata);
279          }
280          if (mock.autoTrigger) {
281            triggers.sendAll();
282          }
283        },
284        sendMessage: (msgBytes: ArrayBufferView) => {
285          if (mock.messageListener) {
286            mock.messageListener(msgBytes);
287          }
288        },
289        finishSend: () => {
290          if (mock.finishSendListener) {
291            mock.finishSendListener();
292          }
293        },
294        cancel: () => {
295          if (mock.cancelListener) {
296            mock.cancelListener();
297          }
298        },
299      };
300    };
301
302    return assignIn(transportConstructor, triggers);
303  }
304}
305
306export interface TriggerableTransport extends grpc.TransportFactory {
307  sendHeaders(): boolean;
308  sendMessages(): boolean;
309  sendTrailers(): boolean;
310  sendAll(): void;
311}