From 249895f979d484184f9d0b7a177b413a41726eb7 Mon Sep 17 00:00:00 2001
From: Keerti Parthasarathy <keertip@google.com>
Date: Wed, 21 Feb 2024 01:04:03 +0000
Subject: [PATCH] Add type hierarchy tests for augmentations for LSP protocol.

Change-Id: Ie6358f4908e988b89afea44b663d03c042346e1c
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/353204
Commit-Queue: Keerti Parthasarathy <keertip@google.com>
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
---
 .../computer_lazy_type_hierarchy.dart         |  5 +-
 .../test/lsp/type_hierarchy_test.dart         | 50 +++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/pkg/analysis_server/lib/src/computer/computer_lazy_type_hierarchy.dart b/pkg/analysis_server/lib/src/computer/computer_lazy_type_hierarchy.dart
index e39d42c764c..42ce0df6949 100644
--- a/pkg/analysis_server/lib/src/computer/computer_lazy_type_hierarchy.dart
+++ b/pkg/analysis_server/lib/src/computer/computer_lazy_type_hierarchy.dart
@@ -130,7 +130,10 @@ class DartLazyTypeHierarchyComputer {
 
     var matches =
         await searchEngine.searchSubtypes(target, SearchEngineCache());
-    return matches.map(toHierarchyItem).toList();
+    return matches
+        .where((match) => !(match.element as InterfaceElement).isAugmentation)
+        .map(toHierarchyItem)
+        .toList();
   }
 
   /// Gets immediate super types for the class/mixin [element].
diff --git a/pkg/analysis_server/test/lsp/type_hierarchy_test.dart b/pkg/analysis_server/test/lsp/type_hierarchy_test.dart
index 4e3de592ed5..b082fb323eb 100644
--- a/pkg/analysis_server/test/lsp/type_hierarchy_test.dart
+++ b/pkg/analysis_server/test/lsp/type_hierarchy_test.dart
@@ -176,6 +176,31 @@ import 'main.dart';
         ]));
   }
 
+  Future<void> test_augment_extends() async {
+    final content = '''
+import augment 'other.dart';
+
+class MyCl^ass1 {}
+[!class /*[1*/C/*1]*/ {}!]
+''';
+    final augmentation = '''
+library augment 'main.dart';
+
+augment class C extends MyClass1 {}
+''';
+    await _fetchSubtypes(content, otherContent: augmentation);
+    expect(
+        subtypes,
+        equals([
+          _isItem(
+            'C',
+            mainFileUri,
+            range: code.ranges[0].range,
+            selectionRange: code.ranges[1].range,
+          ),
+        ]));
+  }
+
   Future<void> test_extends() async {
     final content = '''
 class MyCla^ss1 {}
@@ -309,6 +334,31 @@ class MyCla^ss2 extends MyClass1 {}
         ]));
   }
 
+  Future<void> test_augment_extends() async {
+    final content = '''
+import augment 'other.dart';
+
+[!class /*[1*/MyClass1/*1]*/ {}!]
+class C^s {}
+''';
+    final augmentation = '''
+library augment 'main.dart';
+
+augment class Cs extends MyClass1 {}
+''';
+    await _fetchSupertypes(content, otherContent: augmentation);
+    expect(
+        supertypes,
+        equals([
+          _isItem(
+            'MyClass1',
+            mainFileUri,
+            range: code.ranges[0].range,
+            selectionRange: code.ranges[1].range,
+          ),
+        ]));
+  }
+
   Future<void> test_extends() async {
     final content = '''
 /*[0*/class /*[1*/MyClass1/*1]*/ {}/*0]*/