Bug fixes for the new task model

R=scheglov@google.com

Review URL: https://codereview.chromium.org//1144393006
diff --git a/pkg/analysis_server/lib/src/analysis_server.dart b/pkg/analysis_server/lib/src/analysis_server.dart
index 5ef0dc4..9d8f6b3 100644
--- a/pkg/analysis_server/lib/src/analysis_server.dart
+++ b/pkg/analysis_server/lib/src/analysis_server.dart
@@ -23,6 +23,8 @@
 import 'package:analysis_server/src/source/optimizing_pub_package_map_provider.dart';
 import 'package:analyzer/file_system/file_system.dart';
 import 'package:analyzer/instrumentation/instrumentation.dart';
+import 'package:analyzer/src/context/cache.dart';
+import 'package:analyzer/src/context/context.dart' as newContext;
 import 'package:analyzer/src/generated/ast.dart';
 import 'package:analyzer/src/generated/element.dart';
 import 'package:analyzer/src/generated/engine.dart';
@@ -31,6 +33,8 @@
 import 'package:analyzer/src/generated/source.dart';
 import 'package:analyzer/src/generated/source_io.dart';
 import 'package:analyzer/src/generated/utilities_general.dart';
+import 'package:analyzer/src/task/dart.dart';
+import 'package:analyzer/task/dart.dart';
 import 'package:plugin/plugin.dart';
 
 typedef void OptionUpdater(AnalysisOptionsImpl options);
@@ -1025,10 +1029,21 @@
   void test_flushResolvedUnit(String file) {
     if (AnalysisEngine.isDartFileName(file)) {
       ContextSourcePair contextSource = getContextSourcePair(file);
-      AnalysisContextImpl context = contextSource.context;
+      AnalysisContext context = contextSource.context;
       Source source = contextSource.source;
-      DartEntry dartEntry = context.getReadableSourceEntryOrNull(source);
-      dartEntry.flushAstStructures();
+      if (context is AnalysisContextImpl) {
+        DartEntry dartEntry = context.getReadableSourceEntryOrNull(source);
+        dartEntry.flushAstStructures();
+      } else if (context is newContext.AnalysisContextImpl) {
+        CacheEntry entry = context.getCacheEntry(source);
+        entry.setState(RESOLVED_UNIT1, CacheState.FLUSHED);
+        entry.setState(RESOLVED_UNIT2, CacheState.FLUSHED);
+        entry.setState(RESOLVED_UNIT3, CacheState.FLUSHED);
+        entry.setState(RESOLVED_UNIT4, CacheState.FLUSHED);
+        entry.setState(RESOLVED_UNIT5, CacheState.FLUSHED);
+        entry.setState(RESOLVED_UNIT6, CacheState.FLUSHED);
+        entry.setState(RESOLVED_UNIT, CacheState.FLUSHED);
+      }
     }
   }
 
diff --git a/pkg/analysis_server/test/services/completion/optype_test.dart b/pkg/analysis_server/test/services/completion/optype_test.dart
index 69107b7..27e32b7 100644
--- a/pkg/analysis_server/test/services/completion/optype_test.dart
+++ b/pkg/analysis_server/test/services/completion/optype_test.dart
@@ -1286,5 +1286,8 @@
   @override
   String get shortName => fullName;
 
+  @override
+  Source get source => this;
+
   noSuchMethod(Invocation invocation) => super.noSuchMethod(invocation);
 }
diff --git a/pkg/analysis_server/test/services/refactoring/move_file_test.dart b/pkg/analysis_server/test/services/refactoring/move_file_test.dart
index a3dd17d..cbae32f 100644
--- a/pkg/analysis_server/test/services/refactoring/move_file_test.dart
+++ b/pkg/analysis_server/test/services/refactoring/move_file_test.dart
@@ -145,14 +145,18 @@
     String pathA = '/project/000/1111/a.dart';
     testFile = '/project/000/1111/22/test.dart';
     addSource(pathA, '''
+library lib;
 part '22/test.dart';
 ''');
-    addTestSource('');
+    addTestSource('''
+part of lib;
+''');
     _performAnalysis();
     // perform refactoring
     _createRefactoring('/project/000/1111/22/new_name.dart');
     await _assertSuccessfulRefactoring();
     assertFileChangeResult(pathA, '''
+library lib;
 part '22/new_name.dart';
 ''');
     assertNoFileChange(testFile);
@@ -163,20 +167,26 @@
     String pathB = '/project/000/b.dart';
     testFile = '/project/000/1111/22/test.dart';
     addSource(pathA, '''
+library lib;
 part '22/test.dart';
 ''');
     addSource(pathB, '''
+library lib;
 part '1111/22/test.dart';
 ''');
-    addTestSource('');
+    addTestSource('''
+part of lib;
+''');
     _performAnalysis();
     // perform refactoring
     _createRefactoring('/project/000/1111/22/new_name.dart');
     await _assertSuccessfulRefactoring();
     assertFileChangeResult(pathA, '''
+library lib;
 part '22/new_name.dart';
 ''');
     assertFileChangeResult(pathB, '''
+library lib;
 part '1111/22/new_name.dart';
 ''');
     assertNoFileChange(testFile);