From 93fb8f81812c5aa5814a12e62d336681b32d80ef Mon Sep 17 00:00:00 2001
From: gagik <me@gagik.co>
Date: Fri, 29 Nov 2024 13:26:31 +0100
Subject: [PATCH] add tests

---
 .../suite/participant/participant.test.ts     | 83 +++++++++++++++++++
 1 file changed, 83 insertions(+)

diff --git a/src/test/suite/participant/participant.test.ts b/src/test/suite/participant/participant.test.ts
index ce89bf105..5e4722c2f 100644
--- a/src/test/suite/participant/participant.test.ts
+++ b/src/test/suite/participant/participant.test.ts
@@ -42,6 +42,8 @@ import {
 } from './participantHelpers';
 import EditDocumentCodeLensProvider from '../../../editors/editDocumentCodeLensProvider';
 import PlaygroundResultProvider from '../../../editors/playgroundResultProvider';
+import { CollectionTreeItem, DatabaseTreeItem } from '../../../explorer';
+import type { SendMessageToParticipantOptions } from '../../../participant/participantTypes';
 
 // The Copilot's model in not available in tests,
 // therefore we need to mock its methods and returning values.
@@ -1802,6 +1804,87 @@ Schema:
       });
     });
 
+    suite('opened from tree view', function () {
+      let sendMessageToParticipantStub: SinonStub<
+        [options: SendMessageToParticipantOptions],
+        Promise<unknown>
+      >;
+
+      beforeEach(function () {
+        sendMessageToParticipantStub = sinon.stub(
+          testParticipantController,
+          'sendMessageToParticipant'
+        );
+      });
+
+      suite('with a database item', function () {
+        const mockDatabaseItem = Object.assign(
+          Object.create(DatabaseTreeItem.prototype),
+          {
+            databaseName: 'testDb',
+          } as DatabaseTreeItem
+        );
+
+        test('opens the chat and sends a message to set database context', async function () {
+          expect(sendMessageToParticipantStub).not.called;
+
+          await testParticipantController.askCopilotFromTreeItem(
+            mockDatabaseItem
+          );
+
+          expect(sendMessageToParticipantStub).has.callCount(2);
+
+          expect(sendMessageToParticipantStub.getCall(0).args).deep.equals([
+            {
+              message: `I want to ask questions about the \`${mockDatabaseItem.databaseName}\` database.`,
+              isNewChat: true,
+            },
+          ]);
+
+          expect(sendMessageToParticipantStub.getCall(1).args).deep.equals([
+            {
+              message: '',
+              isPartialQuery: true,
+            },
+          ]);
+        });
+      });
+
+      suite('with a collection item', function () {
+        const mockCollectionItem = Object.assign(
+          Object.create(CollectionTreeItem.prototype),
+          {
+            databaseName: 'testDb',
+            collectionName: 'testColl',
+          } as CollectionTreeItem
+        );
+
+        test('opens the chat and sends a message to set database and collection context', async function () {
+          expect(sendMessageToParticipantStub).not.called;
+
+          await testParticipantController.askCopilotFromTreeItem(
+            mockCollectionItem
+          );
+
+          expect(sendMessageToParticipantStub).has.callCount(2);
+
+          expect(sendMessageToParticipantStub.getCall(0).args).deep.equals([
+            {
+              message: `I want to ask questions about the \`${mockCollectionItem.databaseName}\` database's \`${mockCollectionItem.collectionName}\` collection.`,
+              isNewChat: true,
+            },
+          ]);
+
+          expect(sendMessageToParticipantStub.getCall(1).args).deep.equals([
+            {
+              message: '',
+              isPartialQuery: true,
+            },
+          ]);
+        });
+      });
+    });
+
     suite('determining the namespace', function () {
       ['query', 'schema'].forEach(function (command) {
         suite(`${command} command`, function () {