From 0fd19ab576e7da05cc4150f718458b8c5e2263af Mon Sep 17 00:00:00 2001
From: Simon Binder <oss@simonbinder.eu>
Date: Sun, 10 Sep 2023 22:46:50 +0200
Subject: [PATCH] Fix drift_dev tests

---
 drift_dev/lib/src/cli/commands/migrate.dart   | 35 +++++++++++++------
 .../dart/custom_row_classes_test.dart         |  4 ++-
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/drift_dev/lib/src/cli/commands/migrate.dart b/drift_dev/lib/src/cli/commands/migrate.dart
index d7c2834cd..2c446bf4b 100644
--- a/drift_dev/lib/src/cli/commands/migrate.dart
+++ b/drift_dev/lib/src/cli/commands/migrate.dart
@@ -4,14 +4,17 @@ import 'package:analyzer/dart/analysis/analysis_context.dart';
 import 'package:analyzer/dart/analysis/analysis_context_collection.dart';
 import 'package:analyzer/dart/analysis/results.dart';
 import 'package:analyzer/dart/ast/ast.dart';
+import 'package:analyzer/dart/ast/syntactic_entity.dart';
 import 'package:analyzer/dart/ast/visitor.dart';
+import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/type.dart';
 import 'package:collection/collection.dart';
 import 'package:drift_dev/src/utils/string_escaper.dart';
 import 'package:io/ansi.dart';
 import 'package:path/path.dart' as p;
 import 'package:source_span/source_span.dart';
-import 'package:sqlparser/sqlparser.dart' hide AnalysisContext, StringLiteral;
+import 'package:sqlparser/sqlparser.dart'
+    hide AnalysisContext, StringLiteral, SyntacticEntity, AstNode;
 import 'package:yaml/yaml.dart';
 
 import '../cli.dart';
@@ -386,24 +389,23 @@ class _Moor2DriftDartRewriter extends GeneralizingAstVisitor<void> {
     _writer.replace(uri.offset, uri.length, asDartLiteral(newImport));
   }
 
-  @override
-  void visitSimpleIdentifier(SimpleIdentifier node) {
+  void _transformIdentifier(
+      SyntacticEntity identifier, String name, Element? element) {
     String? newIdentifier;
 
-    if (node.name == 'FlutterQueryExecutor') {
+    if (name == 'FlutterQueryExecutor') {
       newIdentifier = 'SqfliteQueryExecutor';
     } else {
-      var element = node.staticElement;
-      if (element == null) {
+      if (element == null && identifier is AstNode) {
         // It looks like left-hand identifiers of assignments don't have a
         // static element, infer from parent.
-        if (node.parent is AssignmentExpression) {
-          element = (node.parent as AssignmentExpression).writeElement;
+        if (identifier.parent is AssignmentExpression) {
+          element = (identifier.parent as AssignmentExpression).writeElement;
         }
-
-        if (element == null) return;
       }
 
+      if (element == null) return;
+
       for (final annotation in element.metadata) {
         final value = annotation.computeConstantValue();
         if (value == null) return;
@@ -423,10 +425,21 @@ class _Moor2DriftDartRewriter extends GeneralizingAstVisitor<void> {
     }
 
     if (newIdentifier != null) {
-      _writer.replace(node.offset, node.length, newIdentifier);
+      _writer.replace(identifier.offset, name.length, newIdentifier);
     }
   }
 
+  @override
+  void visitSimpleIdentifier(SimpleIdentifier node) {
+    _transformIdentifier(node, node.name, node.staticElement);
+  }
+
+  @override
+  void visitNamedType(NamedType node) {
+    _transformIdentifier(node.name2, node.name2.lexeme, node.element);
+    super.visitNamedType(node);
+  }
+
   @override
   void visitAnnotation(Annotation node) {
     // Replace @UseMoor with @DriftDatabase first
diff --git a/drift_dev/test/analysis/resolver/dart/custom_row_classes_test.dart b/drift_dev/test/analysis/resolver/dart/custom_row_classes_test.dart
index c09811c84..3b657e15d 100644
--- a/drift_dev/test/analysis/resolver/dart/custom_row_classes_test.dart
+++ b/drift_dev/test/analysis/resolver/dart/custom_row_classes_test.dart
@@ -525,7 +525,9 @@ class FooData {
           'a|lib/a.dart': '''
 import 'package:drift/drift.dart';
 
-@UseRowClass((id: int, name: String, birthday: DateTime))
+typedef MyRecord = ({int id, String name, DateTime birthday});
+
+@UseRowClass(MyRecord)
 class Users extends Table {
   IntColumn get id => integer().autoIncrement()();
   TextColumn get name => text()();