From 856af0b1b010bfcab50400b232ffa240c40b07bd Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 15 Nov 2024 14:46:57 +0100
Subject: [PATCH] Add change and tests

---
 src/editors/playgroundController.ts           |  9 +-
 .../editors/playgroundController.test.ts      | 83 +++++++++++++++++++
 2 files changed, 89 insertions(+), 3 deletions(-)

diff --git a/src/editors/playgroundController.ts b/src/editors/playgroundController.ts
index 0e2a6914b..cd69213c4 100644
--- a/src/editors/playgroundController.ts
+++ b/src/editors/playgroundController.ts
@@ -583,10 +583,13 @@ export default class PlaygroundController {
       .get('confirmRunCopilotCode');
 
     if (!this._connectionController.isCurrentlyConnected()) {
-      // TODO(VSCODE-618): Prompt user to connect when clicked.
-      void vscode.window.showErrorMessage(connectBeforeRunningMessage);
+      const successfullyConnected =
+        await this._connectionController.changeActiveConnection();
 
-      return false;
+      if (!successfullyConnected) {
+        void vscode.window.showErrorMessage(connectBeforeRunningMessage);
+        return false;
+      }
     }
 
     if (shouldConfirmRunCopilotCode === true) {
diff --git a/src/test/suite/editors/playgroundController.test.ts b/src/test/suite/editors/playgroundController.test.ts
index 14889416f..57c8ca63c 100644
--- a/src/test/suite/editors/playgroundController.test.ts
+++ b/src/test/suite/editors/playgroundController.test.ts
@@ -270,6 +270,60 @@ suite('Playground Controller Test Suite', function () {
           expectedMessage
         );
       });
+
+      suite('running code from the participant', function () {
+        beforeEach(function () {
+          sinon
+            .stub(testPlaygroundController, '_evaluateWithCancelModal')
+            .resolves({ result: '123' } as any);
+          sinon.stub(testPlaygroundController, '_openInResultPane').resolves();
+
+          showInformationMessageStub.resolves('Yes');
+        });
+
+        afterEach(() => sinon.restore());
+
+        test('prompts to connect to a database and succeeds with selection', async () => {
+          const changeActiveConnectionStub = sinon.stub(
+            testPlaygroundController._connectionController,
+            'changeActiveConnection'
+          );
+          // Mocks the user selecting a connection.
+          changeActiveConnectionStub.resolves(true);
+
+          const result = await testPlaygroundController.evaluateParticipantCode(
+            'console.log("test");'
+          );
+
+          expect(showErrorMessageStub.notCalled).is.true;
+
+          expect(changeActiveConnectionStub.calledOnce).is.true;
+
+          expect(result).is.true;
+        });
+
+        test('prompts to connect to a database and errors if not selected', async () => {
+          const changeActiveConnectionStub = sinon.stub(
+            testPlaygroundController._connectionController,
+            'changeActiveConnection'
+          );
+          // Mocks the user selecting a connection.
+          changeActiveConnectionStub.resolves(false);
+
+          const result = await testPlaygroundController.evaluateParticipantCode(
+            'console.log("test");'
+          );
+
+          const expectedMessage =
+            'Please connect to a database before running a playground.';
+          await testPlaygroundController.runAllOrSelectedPlaygroundBlocks();
+          expect(showErrorMessageStub.firstCall.args[0]).to.be.equal(
+            expectedMessage
+          );
+
+          expect(result).is.false;
+        });
+      });
     });
 
     suite('user is connected', () => {
@@ -431,6 +485,35 @@ suite('Playground Controller Test Suite', function () {
           expect(result).to.be.false;
         });
       });
+
+      suite('running code from the participant', function () {
+        beforeEach(function () {
+          sinon
+            .stub(testPlaygroundController, '_evaluateWithCancelModal')
+            .resolves({ result: '123' } as any);
+          sinon.stub(testPlaygroundController, '_openInResultPane').resolves();
+
+          showInformationMessageStub.resolves('Yes');
+        });
+
+        afterEach(() => sinon.restore());
+
+        test('does not prompt to connect to the database', async () => {
+          const changeActiveConnectionStub = sinon.stub(
+            testPlaygroundController._connectionController,
+            'changeActiveConnection'
+          );
+          const result = await testPlaygroundController.evaluateParticipantCode(
+            'console.log("test");'
+          );
+
+          expect(showErrorMessageStub.notCalled).is.true;
+
+          expect(changeActiveConnectionStub.notCalled).is.true;
+
+          expect(result).is.true;
+        });
+      });
     });
   });
 });