Listen for context add/remove and unit invalidation to update the index.

pkg/analysis_server/test/services/index2/index2_test.dart

Index: pkg/analysis_server/test/services/index2/index2_test.dart
diff --git a/pkg/analysis_server/test/services/index2/index2_test.dart b/pkg/analysis_server/test/services/index2/index2_test.dart
index a23d0c4fa84527293067e6d96b167dc1a605bc60..f65139ab2433f47288cff5750f3fc1f72f7c4ef6 100644
--- a/pkg/analysis_server/test/services/index2/index2_test.dart
+++ b/pkg/analysis_server/test/services/index2/index2_test.dart
@@ -197,12 +197,38 @@ class A {}
     RegExp regExp = new RegExp(r'^A$');
     expect(await index.getDefinedNames(regExp, IndexNameKind.topLevel),
         hasLength(1));
-    // remove the context - no
+    // remove the context - no top-level declarations
     index.removeContext(context);
     expect(
         await index.getDefinedNames(regExp, IndexNameKind.topLevel), isEmpty);
   }
 
+  test_removeUnit() async {
+    RegExp regExp = new RegExp(r'^[AB]$');
+    Source sourceA = addSource('/a.dart', 'class A {}');
+    Source sourceB = addSource('/b.dart', 'class B {}');
+    CompilationUnit unitA = resolveLibraryUnit(sourceA);
+    CompilationUnit unitB = resolveLibraryUnit(sourceB);
+    index.indexUnit(unitA);
+    index.indexUnit(unitB);
+    {
+      List<Location> locations =
+          await index.getDefinedNames(regExp, IndexNameKind.topLevel);
+      expect(locations, hasLength(2));
+      expect(locations.map((l) => l.libraryUri),
+          unorderedEquals([sourceA.uri.toString(), sourceB.uri.toString()]));
+    }
+    // remove a.dart - no a.dart location
+    index.removeUnit(context, sourceA, sourceA);
+    {
+      List<Location> locations =
+          await index.getDefinedNames(regExp, IndexNameKind.topLevel);
+      expect(locations, hasLength(1));
+      expect(locations.map((l) => l.libraryUri),
+          unorderedEquals([sourceB.uri.toString()]));
+    }
+  }
+
   /**
    * Assert that the given list of [locations] has a [Location] corresponding
    * to the [element].