add unit tests

Author: bogdantruta-firebolt

src/statement/stream.ts

@@ -18,9 +18,6 @@ export class StreamStatement {
   }
 
   async streamResult() {
-    if (!this.response.body) {
-      throw new Error("Response body is not readable");
-    }
     return {
       data: new ServerSideStream(this.response, this.executeQueryOptions)
     };

test/unit/v2/connection.test.ts

@@ -400,12 +400,14 @@ describe("Connection V2", () => {
     setupMockServer(server);
     server.use(
       rest.post(
-        `https://some_system_engine.com/${QUERY_URL}?output_format=JSON_Compact`,
+        `https://some_system_engine.com/${QUERY_URL}`,
         async (req, res, ctx) => {
           const body = await req.text();
+          const urlParams = Object.fromEntries(req.url.searchParams.entries());
           if (
             body.includes("fb_GetAsyncStatus") &&
-            body.includes("async_query_token")
+            body.includes("async_query_token") &&
+            urlParams["output_format"] === "JSON_Compact"
           ) {
             return res(
               ctx.json({
@@ -445,4 +447,226 @@ describe("Connection V2", () => {
     await new Promise(resolve => setTimeout(resolve, 100)); // somehow we need it to wait for the flag switch
     expect(cancelQueryExecuted).toBe(true);
   });
+
+  it("streaming works as expected", async () => {
+    const firebolt = Firebolt({
+      apiEndpoint
+    });
+    const jsonLines = [
+      JSON.stringify({
+        message_type: "START",
+        result_columns: [
+          {
+            name: "?column?",
+            type: "integer"
+          }
+        ]
+      }),
+      JSON.stringify({
+        message_type: "DATA",
+        data: [[1], [1]]
+      }),
+      JSON.stringify({
+        message_type: "FINISH_SUCCESSFULLY"
+      })
+    ].join("\n");
+
+    setupMockServer(server);
+    server.use(
+      rest.post(
+        `https://some_system_engine.com/${QUERY_URL}`,
+        async (req, res, ctx) => {
+          const body = await req.text();
+          const urlParams = Object.fromEntries(req.url.searchParams.entries());
+          if (
+            body.includes("select") &&
+            body.includes("generate_series") &&
+            urlParams["output_format"] === "JSONLines_Compact"
+          ) {
+            return res(ctx.body(jsonLines));
+          }
+        }
+      )
+    );
+
+    const connectionParams: ConnectionOptions = {
+      auth: {
+        client_id: "dummy",
+        client_secret: "dummy"
+      },
+      account: "my_account"
+    };
+
+    const connection = await firebolt.connect(connectionParams);
+    const streamStatement = await connection.executeStream(
+      "select 1 from generate_series(1, 2))"
+    );
+    let rowCount = 0;
+    const { data } = await streamStatement.streamResult();
+    data
+      .on("meta", meta => {
+        expect(meta).toEqual([
+          {
+            name: "?column?",
+            type: "integer"
+          }
+        ]);
+      })
+      .on("data", row => {
+        expect(row).toEqual([1]);
+        rowCount++;
+      })
+      .on("end", () => {
+        expect(rowCount).toBe(2);
+      });
+  });
+
+  it("streaming with normalization works as expected", async () => {
+    const firebolt = Firebolt({
+      apiEndpoint
+    });
+    const jsonLines = [
+      JSON.stringify({
+        message_type: "START",
+        result_columns: [
+          {
+            name: "?column?",
+            type: "integer"
+          }
+        ]
+      }),
+      JSON.stringify({
+        message_type: "DATA",
+        data: [[1], [1]]
+      }),
+      JSON.stringify({
+        message_type: "FINISH_SUCCESSFULLY"
+      })
+    ].join("\n");
+
+    setupMockServer(server);
+    server.use(
+      rest.post(
+        `https://some_system_engine.com/${QUERY_URL}`,
+        async (req, res, ctx) => {
+          const body = await req.text();
+          const urlParams = Object.fromEntries(req.url.searchParams.entries());
+          if (
+            body.includes("select") &&
+            body.includes("generate_series") &&
+            urlParams["output_format"] === "JSONLines_Compact"
+          ) {
+            return res(ctx.body(jsonLines));
+          }
+        }
+      )
+    );
+
+    const connectionParams: ConnectionOptions = {
+      auth: {
+        client_id: "dummy",
+        client_secret: "dummy"
+      },
+      account: "my_account"
+    };
+
+    const connection = await firebolt.connect(connectionParams);
+    const streamStatement = await connection.executeStream(
+      "select 1 from generate_series(1, 2))",
+      {
+        response: {
+          normalizeData: true
+        }
+      }
+    );
+    let rowCount = 0;
+    const { data } = await streamStatement.streamResult();
+    data
+      .on("meta", meta => {
+        expect(meta).toEqual([
+          {
+            name: "?column?",
+            type: "integer"
+          }
+        ]);
+      })
+      .on("data", row => {
+        expect(row).toEqual({ "?column?": 1 });
+        rowCount++;
+      })
+      .on("end", () => {
+        expect(rowCount).toBe(2);
+      });
+  });
+
+  it("streaming fails with error", async () => {
+    const firebolt = Firebolt({
+      apiEndpoint
+    });
+    const jsonLines = [
+      JSON.stringify({
+        message_type: "START",
+        result_columns: [
+          {
+            name: "?column?",
+            type: "integer"
+          }
+        ]
+      }),
+      JSON.stringify({
+        message_type: "FINISH_WITH_ERROR"
+      })
+    ].join("\n");
+
+    setupMockServer(server);
+    server.use(
+      rest.post(
+        `https://some_system_engine.com/${QUERY_URL}`,
+        async (req, res, ctx) => {
+          const body = await req.text();
+          const urlParams = Object.fromEntries(req.url.searchParams.entries());
+          if (
+            body.includes("select") &&
+            body.includes("generate_series") &&
+            urlParams["output_format"] === "JSONLines_Compact"
+          ) {
+            return res(ctx.body(jsonLines));
+          }
+        }
+      )
+    );
+
+    const connectionParams: ConnectionOptions = {
+      auth: {
+        client_id: "dummy",
+        client_secret: "dummy"
+      },
+      account: "my_account"
+    };
+
+    const connection = await firebolt.connect(connectionParams);
+    const streamStatement = await connection.executeStream(
+      "select 1 from generate_series(1, 2))"
+    );
+    const { data } = await streamStatement.streamResult();
+    data
+      .on("meta", meta => {
+        expect(meta).toEqual([
+          {
+            name: "?column?",
+            type: "integer"
+          }
+        ]);
+      })
+      .on("data", row => {
+        fail('"Data should not be emitted"');
+      })
+      .on("error", error => {
+        expect(error).toEqual(
+          new Error(
+            'Result encountered an error: {"message_type":"FINISH_WITH_ERROR"}'
+          )
+        );
+      });
+  });
 });