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

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

 // Copyright (c) 2016, the Dart project authors.  Please see the AUTHORS file
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
 import 'package:analysis_server/src/services/index2/index2.dart';
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/src/generated/source.dart';
 import 'package:analyzer/src/summary/idl.dart';
 import 'package:test_reflective_loader/test_reflective_loader.dart';
@@ skipping 179 matching lines @@
     index.indexUnit(testUnit);
   }
 
   test_removeContext() async {
     _indexTestUnit('''
 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].
    */
   void _assertHasDefinedName(List<Location> locations, Element element) {
     String libraryUri = element.library.source.uri.toString();
     String unitUri = element.source.uri.toString();
     for (Location location in locations) {
       if (location.libraryUri == libraryUri &&
           location.unitUri == unitUri &&
@@ skipping 11 matching lines @@
     index.indexUnit(testUnit);
   }
 
   Source _indexUnit(String path, String code) {
     Source source = addSource(path, code);
     CompilationUnit unit = resolveLibraryUnit(source);
     index.indexUnit(unit);
     return source;
   }
 }