Fix an assertion to support migration of part files

I'm not sure where the test for this should go. I found it while
migrating vector_math.dart.

Change-Id: I0c1c5ba91375b2a36540f4c270cef6dc9558a745
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/124103
Commit-Queue: Brian Wilkerson <brianwilkerson@google.com>
Reviewed-by: Paul Berry <paulberry@google.com>

Author: bwilkerson

pkg/nnbd_migration/lib/src/node_builder.dart

@@ -156,7 +156,7 @@ class NodeBuilder extends GeneralizingAstVisitor<DecoratedType>
 
   @override
   DecoratedType visitCompilationUnit(CompilationUnit node) {
-    _graph.migrating(source);
+    _graph.migrating(node.declaredElement.library.source);
     return super.visitCompilationUnit(node);
   }
 

pkg/nnbd_migration/test/api_test.dart

@@ -38,7 +38,7 @@ abstract class _ProvisionalApiTestBase extends AbstractContextTest {
   Future<void> _checkMultipleFileChanges(
       Map<String, String> input, Map<String, String> expectedOutput) async {
     for (var path in input.keys) {
-      newFile(path, content: input[path]);
+      driver.getFileSync(newFile(path, content: input[path]).path);
     }
     var listener = new TestMigrationListener();
     var migration =
@@ -1847,6 +1847,32 @@ bool f(a) => a is List<int?>;
     await _checkSingleFileChanges(content, expected);
   }
 
+  test_libraryWithParts() async {
+    var root = '/home/test/lib';
+    var path1 = convertPath('$root/lib.dart');
+    var file1 = '''
+part 'src/foo/part.dart';
+''';
+    var expected1 = '''
+part 'src/foo/part.dart';
+''';
+    var path2 = convertPath('$root/src/foo/part.dart');
+    var file2 = '''
+part of '../../lib.dart';
+class C {
+  static void m(C c) {}
+}
+''';
+    var expected2 = '''
+part of '../../lib.dart';
+class C {
+  static void m(C c) {}
+}
+''';
+    await _checkMultipleFileChanges(
+        {path2: file2, path1: file1}, {path1: expected1, path2: expected2});
+  }
+
   test_local_function() async {
     var content = '''
 int f(int i) {