Add tests and fix a folding range bug

src/languageserver/handlers/languageHandlers.ts

@@ -30,7 +30,8 @@ export class LanguageHandlers {
   private languageService: LanguageService;
   private yamlSettings: SettingsState;
   private validationHandler: ValidationHandler;
-  private pendingLimitExceededWarnings: { [uri: string]: { features: { [name: string]: string }; timeout?: NodeJS.Timeout } };
+
+  pendingLimitExceededWarnings: { [uri: string]: { features: { [name: string]: string }; timeout?: NodeJS.Timeout } };
 
   constructor(
     private readonly connection: Connection,

src/languageservice/services/yamlFolding.ts

@@ -38,7 +38,7 @@ export function getFoldingRanges(document: TextDocument, context: FoldingRangesC
     context.onRangeLimitExceeded(document.uri);
   }
 
-  return result.slice(0, context.rangeLimit - 1);
+  return result.slice(0, context.rangeLimit);
 }
 
 function creteNormalizedFolding(document: TextDocument, node: ASTNode): FoldingRange {

test/documentSymbols.test.ts

@@ -19,13 +19,29 @@ describe('Document Symbols Tests', () => {
   let languageHandler: LanguageHandlers;
   let yamlSettings: SettingsState;
 
+  const limitContent = `
+    a: [1, 2, 3]
+    b: [4, 5, 6]
+  `;
+
   before(() => {
     const languageSettingsSetup = new ServiceSetup();
     const { languageHandler: langHandler, yamlSettings: settings } = setupLanguageService(languageSettingsSetup.languageSettings);
     languageHandler = langHandler;
     yamlSettings = settings;
   });
 
+  afterEach(() => {
+    yamlSettings.maxItemsComputed = 5000;
+  });
+
+  function assertLimitWarning() {
+    const warnings = languageHandler.pendingLimitExceededWarnings;
+    assert.deepEqual(Object.keys(warnings), [TEST_URI]);
+    assert.deepEqual(warnings[TEST_URI].features, { 'document symbols': 'document symbols' });
+    assert(warnings[TEST_URI].timeout);
+  }
+
   describe('Document Symbols Tests (Non Hierarchical)', function () {
     function parseNonHierarchicalSetup(content: string): SymbolInformation[] | DocumentSymbol[] {
       const testTextDocument = setupTextDocument(content);
@@ -121,6 +137,16 @@ describe('Document Symbols Tests', () => {
       assert.deepEqual(symbols[0], createExpectedSymbolInformation('analytics', 17, '', TEST_URI, 1, 0, 1, 15));
       assert.deepEqual(symbols[1], createExpectedSymbolInformation('json', 15, '', TEST_URI, 4, 0, 4, 10));
     });
+
+    it('Document symbols with a limit', () => {
+      yamlSettings.maxItemsComputed = 1;
+
+      const symbols = parseNonHierarchicalSetup(limitContent);
+      assert.equal(symbols.length, 1);
+      assert.deepEqual(symbols[0], createExpectedSymbolInformation('a', SymbolKind.Array, '', TEST_URI, 1, 4, 1, 16));
+
+      assertLimitWarning();
+    });
   });
 
   describe('Document Symbols Tests (Hierarchical)', function () {
@@ -275,5 +301,21 @@ describe('Document Symbols Tests', () => {
         createExpectedDocumentSymbol('conditions', SymbolKind.Module, 6, 12, 10, 28, 6, 12, 6, 22, [root2])
       );
     });
+
+    it('Document symbols with a limit', () => {
+      yamlSettings.maxItemsComputed = 3;
+
+      const symbols = parseHierarchicalSetup(limitContent) as DocumentSymbol[];
+      assert.equal(symbols.length, 2);
+      assert.equal(symbols[0].children.length, 1);
+      assert.equal(symbols[1].children.length, 0);
+
+      const el = createExpectedDocumentSymbolNoDetail('0', SymbolKind.Number, 1, 8, 1, 9, 1, 8, 1, 9, []);
+      const root = createExpectedDocumentSymbol('a', SymbolKind.Array, 1, 4, 1, 16, 1, 4, 1, 5, [el]);
+
+      assert.deepEqual(symbols[0], root);
+
+      assertLimitWarning();
+    });
   });
 });

test/yamlFolding.test.ts

@@ -7,7 +7,7 @@ import { expect } from 'chai';
 import { FoldingRange } from 'vscode-languageserver';
 import { getFoldingRanges } from '../src/languageservice/services/yamlFolding';
 import { FoldingRangesContext } from '../src/languageservice/yamlTypes';
-import { setupTextDocument } from './utils/testHelper';
+import { setupTextDocument, TEST_URI } from './utils/testHelper';
 
 const context: FoldingRangesContext = { rangeLimit: 10_0000 };
 
@@ -100,4 +100,31 @@ SecondDict:
     const ranges = getFoldingRanges(doc, context);
     expect(ranges).to.deep.include.members([FoldingRange.create(1, 4, 2, 15)]);
   });
+
+  it('should respect range limits', () => {
+    const yaml = `
+      a:
+        - 1
+      b:
+        - 2
+    `
+
+    const warnings = [];
+
+    const doc = setupTextDocument(yaml);
+
+    const unlimitedRanges = getFoldingRanges(doc, {
+      rangeLimit: 10,
+      onRangeLimitExceeded: (uri) => warnings.push(uri),
+    });
+    expect(unlimitedRanges.length).to.equal(2);
+    expect(warnings).to.be.empty;
+
+    const limitedRanges = getFoldingRanges(doc, {
+      rangeLimit: 1,
+      onRangeLimitExceeded: (uri) => warnings.push(uri),
+    });
+    expect(limitedRanges.length).to.equal(1);
+    expect(warnings).to.deep.equal([TEST_URI]);
+  });
 });