feat: provide shell methods completions after getCollection VSCODE-390 (#498)

Author: alenakhineika

src/language/visitor.ts

@@ -575,7 +575,7 @@ export class Visitor {
     }
   }
 
-  _checkIsCollectionSymbol(node: babel.types.MemberExpression): void {
+  _checkIsCollectionMemberExpression(node: babel.types.MemberExpression): void {
     if (
       node.object.type === 'MemberExpression' &&
       node.object.object.type === 'Identifier' &&
@@ -586,4 +586,24 @@ export class Visitor {
       this._state.isCollectionSymbol = true;
     }
   }
+
+  _checkIsCollectionCallExpression(node: babel.types.MemberExpression): void {
+    if (
+      node.object.type === 'CallExpression' &&
+      node.object.callee.type === 'MemberExpression' &&
+      node.object.callee.object.type === 'Identifier' &&
+      node.object.callee.object.name === 'db' &&
+      node.object.callee.property.type === 'Identifier' &&
+      node.object.callee.property.name === 'getCollection' &&
+      node.property.type === 'Identifier' &&
+      node.property.name.includes(PLACEHOLDER)
+    ) {
+      this._state.isCollectionSymbol = true;
+    }
+  }
+
+  _checkIsCollectionSymbol(node: babel.types.MemberExpression): void {
+    this._checkIsCollectionMemberExpression(node);
+    this._checkIsCollectionCallExpression(node);
+  }
 }

src/test/suite/language/mongoDBService.test.ts

@@ -151,7 +151,7 @@ suite('MongoDBService Test Suite', () => {
       expect(completion).to.have.property('kind', CompletionItemKind.Method);
     });
 
-    test('provide shell collection methods completion if collection name is computed property', async () => {
+    test('provide shell collection methods completion for a collection name in a bracket notation', async () => {
       const result = await testMongoDBService.provideCompletionItems(
         ['use("test");', 'db["test"].'].join('\n'),
         { line: 1, character: 11 }
@@ -163,6 +163,18 @@ suite('MongoDBService Test Suite', () => {
       expect(completion).to.have.property('kind', CompletionItemKind.Method);
     });
 
+    test('provide shell collection methods completion for a collection name in getCollection', async () => {
+      const result = await testMongoDBService.provideCompletionItems(
+        ['use("test");', 'db.getCollection("test").'].join('\n'),
+        { line: 1, character: 41 }
+      );
+      const completion = result.find(
+        (item: CompletionItem) => item.label === 'find'
+      );
+
+      expect(completion).to.have.property('kind', CompletionItemKind.Method);
+    });
+
     test('provide shell collection methods completion if single quotes', async () => {
       const result = await testMongoDBService.provideCompletionItems(
         ["use('test');", "db['test']."].join('\n'),