diff --git a/moor_generator/lib/src/analyzer/runner/steps.dart b/moor_generator/lib/src/analyzer/runner/steps.dart
index 44ae604b..b3c16b82 100644
--- a/moor_generator/lib/src/analyzer/runner/steps.dart
+++ b/moor_generator/lib/src/analyzer/runner/steps.dart
@@ -33,7 +33,9 @@ abstract class Step {
 
   String get path => file.uri.path;
 
-  Step(this.task, this.file);
+  Step(this.task, this.file)
+      : assert(task != null),
+        assert(file != null);
 
   void reportError(MoorError error) =>
       errors.report(error..wasDuringParsing = isParsing);
diff --git a/moor_generator/test/analyzer/dart/dart_test.dart b/moor_generator/test/analyzer/dart/dart_test.dart
index 9b1223e5..f0879bd6 100644
--- a/moor_generator/test/analyzer/dart/dart_test.dart
+++ b/moor_generator/test/analyzer/dart/dart_test.dart
@@ -5,6 +5,7 @@ import 'package:analyzer/dart/element/element.dart';
 import 'package:build/build.dart';
 import 'package:moor_generator/src/analyzer/dart/parser.dart';
 import 'package:moor_generator/src/analyzer/runner/steps.dart';
+import 'package:moor_generator/src/analyzer/session.dart';
 import 'package:test/test.dart';
 
 import '../../utils/test_backend.dart';
@@ -24,10 +25,13 @@ void main() {
     });
 
     final input = Uri.parse('package:test_lib/main.dart');
+    final session = MoorSession(backend);
     final backendTask = backend.startTask(input);
+    final task = session.startTask(backendTask);
 
     final library = await backendTask.resolveDart(input);
-    final parser = MoorDartParser(ParseDartStep(null, null, library));
+    final parser = MoorDartParser(
+        ParseDartStep(task, session.registerFile(input), library));
 
     Future<MethodDeclaration> _loadDeclaration(Element element) async {
       final declaration = await parser.loadElementDeclaration(element);
diff --git a/moor_generator/test/analyzer/dart/table_parser_test.dart b/moor_generator/test/analyzer/dart/table_parser_test.dart
index 82d91c32..0e9094bf 100644
--- a/moor_generator/test/analyzer/dart/table_parser_test.dart
+++ b/moor_generator/test/analyzer/dart/table_parser_test.dart
@@ -4,6 +4,7 @@ import 'package:analyzer/dart/element/element.dart';
 import 'package:build/build.dart';
 import 'package:moor_generator/src/analyzer/dart/parser.dart';
 import 'package:moor_generator/src/analyzer/runner/steps.dart';
+import 'package:moor_generator/src/analyzer/session.dart';
 import 'package:moor_generator/src/model/specified_column.dart';
 import 'package:moor_generator/src/model/specified_table.dart';
 import 'package:test/test.dart';
@@ -75,8 +76,12 @@ void main() {
   setUp(() async {
     final uri = Uri.parse('package:test_lib/main.dart');
     final task = backend.startTask(uri);
+    final session = MoorSession(backend);
 
-    dartStep = ParseDartStep(null, null, await task.resolveDart(uri));
+    final moorTask = session.startTask(task);
+    final file = session.registerFile(uri);
+
+    dartStep = ParseDartStep(moorTask, file, await task.resolveDart(uri));
     parser = MoorDartParser(dartStep);
   });
 
diff --git a/moor_generator/test/analyzer/moor/moor_parser_test.dart b/moor_generator/test/analyzer/moor/moor_parser_test.dart
index 45e70815..a89b685a 100644
--- a/moor_generator/test/analyzer/moor/moor_parser_test.dart
+++ b/moor_generator/test/analyzer/moor/moor_parser_test.dart
@@ -1,6 +1,10 @@
+import 'package:build/build.dart';
 import 'package:moor_generator/src/analyzer/runner/steps.dart';
+import 'package:moor_generator/src/analyzer/session.dart';
 import 'package:test/test.dart';
 
+import '../../utils/test_backend.dart';
+
 void main() {
   const content = '''
 import 'package:my_package/some_file.dart';
@@ -18,7 +22,13 @@ usersWithLongName: SELECT * FROM users WHERE LENGTH(name) > 25
   ''';
 
   test('parses standalone .moor files', () async {
-    final parseStep = ParseMoorStep(null, null, content);
+    final asset = AssetId.parse('foo|bar.moor');
+    final backend = TestBackend({asset: content});
+    final session = MoorSession(backend);
+    final task = session.startTask(backend.startTask(asset.uri));
+    final file = session.registerFile(asset.uri);
+
+    final parseStep = ParseMoorStep(task, file, content);
     final result = await parseStep.parseFile();
 
     expect(parseStep.errors.errors, isEmpty);
@@ -33,5 +43,7 @@ usersWithLongName: SELECT * FROM users WHERE LENGTH(name) > 25
         ['int', 'String', 'bool', 'DateTime', 'int']);
     expect(table.columns.map((c) => c.getJsonKey()),
         ['id', 'name', 'field', 'another', 'myJsonKey']);
+
+    backend.finish();
   });
 }