From ecdfa216da12fd58c6042ba2cfb245c92001d02c Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Mon, 15 Apr 2024 12:29:42 +0200
Subject: [PATCH] fix: throw an error if test data is not found (#343)

---
 src/testTree.ts       | 1 +
 src/testTreeData.ts   | 5 ++++-
 test/TestData.test.ts | 6 +++++-
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/testTree.ts b/src/testTree.ts
index 9bdc38a..175f37b 100644
--- a/src/testTree.ts
+++ b/src/testTree.ts
@@ -91,6 +91,7 @@ export class TestTree extends vscode.Disposable {
       range: undefined,
       error: undefined,
     }
+    TestFolder.register(item)
     this.folderItems.set(id, item)
     return item
   }
diff --git a/src/testTreeData.ts b/src/testTreeData.ts
index fc18529..beac029 100644
--- a/src/testTreeData.ts
+++ b/src/testTreeData.ts
@@ -6,7 +6,10 @@ export type TestData = TestFolder | TestFile | TestCase | TestSuite
 const WEAKMAP_TEST_DATA = new WeakMap<vscode.TestItem, TestData>()
 
 export function getTestData(item: vscode.TestItem): TestData {
-  return WEAKMAP_TEST_DATA.get(item)!
+  const data = WEAKMAP_TEST_DATA.get(item)
+  if (!data)
+    throw new Error(`Test data not found for "${item.label}". This is a bug in Vitest extension. Please report it to https://github.com/vitest-dev/vscode`)
+  return data
 }
 
 export function addTestData<T extends TestData>(item: vscode.TestItem, data: T): T {
diff --git a/test/TestData.test.ts b/test/TestData.test.ts
index a42c565..473bf76 100644
--- a/test/TestData.test.ts
+++ b/test/TestData.test.ts
@@ -1,7 +1,7 @@
 import * as path from 'node:path'
 import * as vscode from 'vscode'
 import { expect } from 'chai'
-import { TestCase, TestFile, TestSuite } from '../src/testTreeData'
+import { TestCase, TestFile, TestSuite, getTestData } from '../src/testTreeData'
 
 describe('TestData', () => {
   const ctrl = vscode.tests.createTestController('mocha', 'Vitest')
@@ -64,5 +64,9 @@ describe('TestData', () => {
       expect(test2.getTestNamePattern()).to.equal('^\\s?describe test 1$')
       expect(test3.getTestNamePattern()).to.equal('^\\s?describe test 2$')
     })
+
+    it('throws an error if data was not set', () => {
+      expect(() => getTestData({ label: 'invalid test' } as any)).to.throw(/Test data not found for "invalid test"/)
+    })
   })
 })
-- 
2.51.2