From 448409536b32d19be1291db7e163e1615d4f9191 Mon Sep 17 00:00:00 2001
From: Philip Carneiro <pcarneiro@kx.com>
Date: Thu, 8 Feb 2024 10:35:27 +0000
Subject: [PATCH] add more tests

---
 src/commands/serverCommand.ts |  6 +----
 test/suite/commands.test.ts   | 47 +++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/commands/serverCommand.ts b/src/commands/serverCommand.ts
index b7c1342b..0fe256e9 100644
--- a/src/commands/serverCommand.ts
+++ b/src/commands/serverCommand.ts
@@ -942,14 +942,10 @@ export function rerunQuery(rerunQueryElement: QueryHistory) {
       context,
       rerunQueryElement.language !== "q",
     );
-  } else if (
-    rerunQueryElement.datasourceType &&
-    rerunQueryElement.isDatasource
-  ) {
+  } else {
     const dsFile = rerunQueryElement.query as DataSourceFiles;
     runDataSource(dsFile);
   }
-  console.log(rerunQueryElement);
 }
 
 export async function loadServerObjects(): Promise<ServerObject[]> {
diff --git a/test/suite/commands.test.ts b/test/suite/commands.test.ts
index 6f198860..abe6722f 100644
--- a/test/suite/commands.test.ts
+++ b/test/suite/commands.test.ts
@@ -43,6 +43,8 @@ import * as dataSourceUtils from "../../src/utils/dataSource";
 import { ExecutionConsole } from "../../src/utils/executionConsole";
 import * as queryUtils from "../../src/utils/queryUtils";
 import { Connection } from "../../src/models/connection";
+import { QueryHistory } from "../../src/models/queryHistory";
+import { ServerType } from "../../src/models/server";
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const dsCmd = require("../../src/commands/dataSourceCommand");
@@ -1689,6 +1691,51 @@ describe("serverCommand", () => {
       );
     });
   });
+
+  describe("rerunQuery", function () {
+    let executeQueryStub, runDataSourceStub: sinon.SinonStub;
+    beforeEach(() => {
+      runDataSourceStub = sinon
+        .stub(dataSourceCommand, "runDataSource")
+        .resolves();
+
+      executeQueryStub = sinon.stub(serverCommand, "executeQuery").resolves();
+    });
+    this.afterEach(() => {
+      sinon.restore();
+    });
+    it("should execute query for non-datasource query", async function () {
+      const rerunQueryElement: QueryHistory = {
+        isDatasource: false,
+        query: "SELECT * FROM table",
+        language: "q",
+        time: "",
+        success: true,
+        connectionName: "",
+        connectionType: ServerType.KDB,
+      };
+
+      serverCommand.rerunQuery(rerunQueryElement);
+      sinon.assert.notCalled(runDataSourceStub);
+    });
+
+    it("should run datasource for datasource query", async function () {
+      const ds = createDefaultDataSourceFile();
+      const rerunQueryElement: QueryHistory = {
+        isDatasource: true,
+        datasourceType: DataSourceTypes.QSQL,
+        query: ds,
+        connectionName: "",
+        connectionType: ServerType.INSIGHTS,
+        time: "",
+        success: false,
+      };
+
+      await serverCommand.rerunQuery(rerunQueryElement);
+
+      sinon.assert.notCalled(executeQueryStub);
+    });
+  });
 });
 
 describe("walkthroughCommand", () => {