AnalysisDriver.getUnitElement() should return a structure with a key.

pkg/analysis_server/test/src/utilities/change_builder_dart_test.dart

 // Copyright (c) 2015, 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 'dart:async';
+
 import 'package:analysis_server/plugin/protocol/protocol.dart';
 import 'package:analysis_server/src/provisional/edit/utilities/change_builder_core.dart';
 import 'package:analysis_server/src/provisional/edit/utilities/change_builder_dart.dart';
 import 'package:analysis_server/src/utilities/change_builder_dart.dart';
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/ast/standard_resolution_map.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/type.dart';
+import 'package:analyzer/src/dart/analysis/driver.dart';
 import 'package:test/test.dart';
 import 'package:test_reflective_loader/test_reflective_loader.dart';
 
 import '../../abstract_context.dart';
 
 main() {
   defineReflectiveSuite(() {
     defineReflectiveTests(DartChangeBuilderImplTest);
     defineReflectiveTests(DartEditBuilderImplTest);
     defineReflectiveTests(DartFileEditBuilderImplTest);
@@ skipping 32 matching lines @@
     SourceFileEdit fileEdit = fileEdits[0];
     expect(fileEdit, isNotNull);
     List<SourceEdit> edits = fileEdit.edits;
     expect(edits, hasLength(1));
     return edits[0];
   }
 
   test_writeClassDeclaration_interfaces() async {
     String path = '/test.dart';
     addSource(path, 'class A {}');
-    DartType typeA = (await driver.getUnitElement(path)).getType('A').type;
+    DartType typeA = await _getType(path, 'A');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(0, (EditBuilder builder) {
         (builder as DartEditBuilder)
             .writeClassDeclaration('C', interfaces: [typeA]);
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(
@@ skipping 28 matching lines @@
         });
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('class C { /**/ }'));
   }
 
   test_writeClassDeclaration_mixins_noSuperclass() async {
     String path = '/test.dart';
     addSource(path, 'class A {}');
-    DartType typeA = (await driver.getUnitElement(path)).getType('A').type;
+    DartType typeA = await _getType(path, 'A');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(0, (EditBuilder builder) {
         (builder as DartEditBuilder)
             .writeClassDeclaration('C', mixins: [typeA]);
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement,
         equalsIgnoringWhitespace('class C extends Object with A { }'));
   }
 
   test_writeClassDeclaration_mixins_superclass() async {
     String path = '/test.dart';
     addSource(path, 'class A {} class B {}');
-    DartType typeA = (await driver.getUnitElement(path)).getType('A').type;
-    DartType typeB = (await driver.getUnitElement(path)).getType('B').type;
+    DartType typeA = await _getType(path, 'A');
+    DartType typeB = await _getType(path, 'B');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(0, (EditBuilder builder) {
         (builder as DartEditBuilder)
             .writeClassDeclaration('C', mixins: [typeB], superclass: typeA);
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement,
@@ skipping 18 matching lines @@
         builder.sourceChange.linkedEditGroups;
     expect(linkedEditGroups, hasLength(1));
     LinkedEditGroup group = linkedEditGroups[0];
     expect(group.length, 1);
     expect(group.positions, hasLength(1));
   }
 
   test_writeClassDeclaration_superclass() async {
     String path = '/test.dart';
     addSource(path, 'class B {}');
-    DartType typeB = (await driver.getUnitElement(path)).getType('B').type;
+    DartType typeB = await _getType(path, 'B');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(0, (EditBuilder builder) {
         (builder as DartEditBuilder)
             .writeClassDeclaration('C', superclass: typeB);
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('class C extends B { }'));
@@ skipping 100 matching lines @@
     expect(group.length, 1);
     expect(group.positions, hasLength(1));
     Position position = group.positions[0];
     expect(position.offset, equals(13));
   }
 
   test_writeFieldDeclaration_type_typeGroupName() async {
     String path = '/test.dart';
     String content = 'class A {} class B {}';
     addSource(path, content);
-    DartType typeA = (await driver.getUnitElement(path)).getType('A').type;
+    DartType typeA = await _getType(path, 'A');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilder)
             .writeFieldDeclaration('f', type: typeA, typeGroupName: 'type');
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('A f;'));
@@ skipping 64 matching lines @@
     expect(group.length, 1);
     expect(group.positions, hasLength(1));
     Position position = group.positions[0];
     expect(position.offset, equals(13));
   }
 
   test_writeGetterDeclaration_returnType() async {
     String path = '/test.dart';
     String content = 'class A {} class B {}';
     addSource(path, content);
-    DartType typeA = (await driver.getUnitElement(path)).getType('A').type;
+    DartType typeA = await _getType(path, 'A');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilder).writeGetterDeclaration('g',
             returnType: typeA, returnTypeGroupName: 'returnType');
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('A get g => null;'));
@@ skipping 10 matching lines @@
 
   test_writeOverrideOfInheritedMember() async {
     String path = '/test.dart';
     String content = '''
 class A {
   A add(A a) => null;
 }
 class B extends A {
 }''';
     addSource(path, content);
-    ClassElement classA = (await driver.getUnitElement(path)).getType('A');
+    ClassElement classA = await _getClassElement(path, 'A');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilder)
             .writeOverrideOfInheritedMember(classA.methods[0]);
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('''
@@ skipping 110 matching lines @@
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('(String s, int i)'));
   }
 
   test_writeParameterSource() async {
     String path = '/test.dart';
     String content = 'class A {}';
     addSource(path, content);
-    DartType typeA = (await driver.getUnitElement(path)).getType('A').type;
+    DartType typeA = await _getType(path, 'A');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilder).writeParameterSource(typeA, 'a');
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('A a'));
   }
@@ skipping 15 matching lines @@
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace(''));
   }
 
   test_writeType_genericType() async {
     String path = '/test.dart';
     String content = 'class A {} class B<E> {}';
     addSource(path, content);
-    InterfaceType typeA = (await driver.getUnitElement(path)).getType('A').type;
-    InterfaceType typeB = (await driver.getUnitElement(path)).getType('B').type;
+    InterfaceType typeA = await _getType(path, 'A');
+    InterfaceType typeB = await _getType(path, 'B');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilder).writeType(typeB.instantiate([typeA]));
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('B<A>'));
   }
 
   test_writeType_groupName() async {
     String path = '/test.dart';
     String content = 'class A {} class B extends A {} class C extends B {}';
     addSource(path, content);
-    DartType typeC = (await driver.getUnitElement(path)).getType('C').type;
+    DartType typeC = await _getType(path, 'C');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilder).writeType(typeC, groupName: 'type');
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('C'));
 
     List<LinkedEditGroup> linkedEditGroups =
         builder.sourceChange.linkedEditGroups;
     expect(linkedEditGroups, hasLength(1));
     LinkedEditGroup group = linkedEditGroups[0];
     expect(group, isNotNull);
   }
 
   test_writeType_groupName_addSupertypeProposals() async {
     String path = '/test.dart';
     String content = 'class A {} class B extends A {} class C extends B {}';
     addSource(path, content);
-    DartType typeC = (await driver.getUnitElement(path)).getType('C').type;
+    DartType typeC = await _getType(path, 'C');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilder)
             .writeType(typeC, addSupertypeProposals: true, groupName: 'type');
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('C'));
@@ skipping 49 matching lines @@
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('var'));
   }
 
   test_writeType_required_notNull() async {
     String path = '/test.dart';
     String content = 'class A {}';
     addSource(path, content);
-    DartType typeA = (await driver.getUnitElement(path)).getType('A').type;
+    DartType typeA = await _getType(path, 'A');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilder).writeType(typeA, required: true);
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('A'));
   }
@@ skipping 10 matching lines @@
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('var'));
   }
 
   test_writeType_simpleType() async {
     String path = '/test.dart';
     String content = 'class A {}';
     addSource(path, content);
-    DartType typeA = (await driver.getUnitElement(path)).getType('A').type;
+    DartType typeA = await _getType(path, 'A');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilder).writeType(typeA);
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('A'));
   }
@@ skipping 10 matching lines @@
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, isEmpty);
   }
 
   test_writeTypes_noPrefix() async {
     String path = '/test.dart';
     String content = 'class A {} class B {}';
     addSource(path, content);
-    DartType typeA = (await driver.getUnitElement(path)).getType('A').type;
-    DartType typeB = (await driver.getUnitElement(path)).getType('B').type;
+    DartType typeA = await _getType(path, 'A');
+    DartType typeB = await _getType(path, 'B');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilderImpl).writeTypes([typeA, typeB]);
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('A, B'));
   }
@@ skipping 10 matching lines @@
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, isEmpty);
   }
 
   test_writeTypes_prefix() async {
     String path = '/test.dart';
     String content = 'class A {} class B {}';
     addSource(path, content);
-    DartType typeA = (await driver.getUnitElement(path)).getType('A').type;
-    DartType typeB = (await driver.getUnitElement(path)).getType('B').type;
+    DartType typeA = await _getType(path, 'A');
+    DartType typeB = await _getType(path, 'B');
 
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, 1, (FileEditBuilder builder) {
       builder.addInsertion(content.length - 1, (EditBuilder builder) {
         (builder as DartEditBuilderImpl)
             .writeTypes([typeA, typeB], prefix: 'implements ');
       });
     });
     SourceEdit edit = getEdit(builder);
     expect(edit.replacement, equalsIgnoringWhitespace('implements A, B'));
   }
+
+  Future<ClassElement> _getClassElement(String path, String name) async {
+    UnitElementResult result = await driver.getUnitElement(path);
+    return result.element.getType(name);
+  }
+
+  Future<DartType> _getType(String path, String name) async {
+    ClassElement classElement = await _getClassElement(path, name);
+    return classElement.type;
+  }
 }
 
 @reflectiveTest
 class DartFileEditBuilderImplTest extends AbstractContextTest {
   @override
   bool get enableNewAnalysisDriver => true;
 
   test_createEditBuilder() async {
     String path = '/test.dart';
     addSource(path, 'library test;');
     int timeStamp = 65;
     DartChangeBuilderImpl builder = new DartChangeBuilderImpl(driver);
     await builder.addFileEdit(path, timeStamp, (FileEditBuilder builder) {
       int offset = 4;
       int length = 5;
       DartEditBuilderImpl editBuilder = (builder as DartFileEditBuilderImpl)
           .createEditBuilder(offset, length);
       expect(editBuilder, new isInstanceOf<DartEditBuilder>());
       SourceEdit sourceEdit = editBuilder.sourceEdit;
       expect(sourceEdit.length, length);
       expect(sourceEdit.offset, offset);
       expect(sourceEdit.replacement, isEmpty);
     });
   }
 }