Fix family args symbol imports

lib/src/code_builders/class_member_copier.dart

@@ -20,38 +20,50 @@ class MemberCopierImpl implements MemberCopier {
 
   @override
   Method? copyGetter(GetterElement element) =>
-      copyGetterImplementation(element)?.let((implementation) {
-        return Method((mb) => mb
-          ..name = element.name
-          ..type = MethodType.getter
-          ..lambda = true
-          ..body = implementation);
+      copyGetterImplementation(element, library)?.let((implementation) {
+        return Method(
+          (mb) => mb
+            ..name = element.name
+            ..type = MethodType.getter
+            ..lambda = true
+            ..body = implementation,
+        );
       });
 
   @override
   Field? copyField(FieldElement element) =>
-      copyFieldImplementation(element)?.let((implementation) {
-        return Field((fb) => fb
-          ..name = element.name
-          ..late = element.isLate
-          ..modifier =
-              element.isFinal ? FieldModifier.final$ : FieldModifier.var$
-          ..assignment = implementation);
+      copyFieldImplementation(element, library)?.let((implementation) {
+        return Field(
+          (fb) => fb
+            ..name = element.name
+            ..late = element.isLate
+            ..modifier = element.isFinal
+                ? FieldModifier.final$
+                : FieldModifier.var$
+            ..assignment = implementation,
+        );
       });
 
   @override
   Method? copyMethod(MethodElement element, bool isOverride) =>
       copyMethodImplementation(element, library)?.let((implementation) {
-        return Method((mb) => mb
-          ..name = element.name
-          ..body = implementation
-          ..annotations
-              .addAll([if (isOverride) CodeExpression(Code("override"))]));
+        return Method(
+          (mb) => mb
+            ..name = element.name
+            ..body = implementation
+            ..annotations.addAll([
+              if (isOverride) CodeExpression(Code("override")),
+            ]),
+        );
       });
 
-  Code? copyGetterImplementation(TypeParameterizedElement element) {
-    final declaration =
-        library.getFragmentDeclaration(element.firstFragment)?.node;
+  Code? copyGetterImplementation(
+    TypeParameterizedElement element,
+    ResolvedLibraryResult resolvedLibrary,
+  ) {
+    final declaration = resolvedLibrary
+        .getFragmentDeclaration(element.firstFragment)
+        ?.node;
     if (declaration == null) {
       // Fields early exit here
       return null;
@@ -61,43 +73,60 @@ class MemberCopierImpl implements MemberCopier {
     }
     final rawBody = declaration.body;
     final bodySource = switch (rawBody) {
-      ExpressionFunctionBody(:final expression) =>
-        CodeExpression(Code(expression.toSource())).code,
-      BlockFunctionBody(:final block) => Code(block.statements
-              .whereType<ReturnStatement>()
-              .firstOrNull
-              ?.toSource() ??
-          "// Found nothing"),
+      ExpressionFunctionBody(:final expression) => CodeExpression(
+        Code(expression.toSource()),
+      ).code,
+      BlockFunctionBody(:final block) => Code(
+        block.statements.whereType<ReturnStatement>().firstOrNull?.toSource() ??
+            "// Found nothing",
+      ),
       _ => Code(rawBody.toSource()),
     };
     return bodySource;
   }
 
-  Code? copyFieldImplementation(FieldElement element) {
-    final library =
-        element.session?.getParsedLibraryByElement(element.library);
-    if (library is! ParsedLibraryResult) {
-      return null;
-    }
-    final declaration =
-        library.getFragmentDeclaration(element.firstFragment)?.node;
+  Code? copyFieldImplementation(
+    FieldElement element,
+    ResolvedLibraryResult resolvedLibrary,
+  ) {
+    final declaration = resolvedLibrary
+        .getFragmentDeclaration(element.firstFragment)
+        ?.node;
     if (declaration == null) {
       return null;
     }
     if (declaration is! VariableDeclaration) {
       return null;
     }
     final rawBody = declaration.initializer;
+    if (rawBody == null) {
+      return Code("");
+    }
+    final allocator = SymbolResolver();
+    declaration.accept(allocator);
+    final references = allocator.consumeReferences();
     return Code.scope((allocate) {
+      String bodySource = rawBody.toSource();
       allocate(
-          refer("Ref", "package:riverpod_annotation/riverpod_annotation.dart"));
-      return rawBody?.toSource() ?? "";
+        refer("Ref", "package:riverpod_annotation/riverpod_annotation.dart"),
+      );
+      for (final reference in references) {
+        bodySource = bodySource.replaceAll(
+          reference.symbol ?? "",
+          allocate(reference),
+        );
+      }
+      return bodySource;
     });
   }
 
-  Code? copyMethodImplementation(MethodElement element, ResolvedLibraryResult resolvedLibrary) {
-    final declaration =
-        resolvedLibrary.getFragmentDeclaration(element.firstFragment)?.node;
+  Code? copyMethodImplementation(
+    MethodElement element,
+    ResolvedLibraryResult resolvedLibrary,
+  ) {
+    final declaration = resolvedLibrary
+        .getFragmentDeclaration(element.firstFragment)
+        ?.node;
     if (declaration == null) {
       return null;
     }
@@ -111,7 +140,10 @@ class MemberCopierImpl implements MemberCopier {
     return Code.scope((allocate) {
       String bodySource = body.toSource();
       for (final reference in references) {
-        bodySource = bodySource.replaceAll(reference.symbol ?? "", allocate(reference));
+        bodySource = bodySource.replaceAll(
+          reference.symbol ?? "",
+          allocate(reference),
+        );
       }
       return bodySource;
     });

test/golden/flumepod/family_notifier/expected.flumepod.dart

@@ -7,17 +7,19 @@ import 'package:riverpod/src/providers/notifier.dart' as _i1;
 import 'package:mockito/mockito.dart' as _i2;
 import 'input.dart' as _i3;
 import 'package:riverpod_annotation/riverpod_annotation.dart' as _i4;
-import 'package:riverpod/src/framework.dart' as _i5;
-import 'package:mockito/src/dummies.dart' as _i6;
-import 'dart:async' as _i7;
+import 'asset:flumepod/test/model.dart' as _i5;
+import 'package:riverpod/src/framework.dart' as _i6;
+import 'model.dart' as _i7;
+import 'package:mockito/src/dummies.dart' as _i8;
+import 'dart:async' as _i9;
 
 // Flumepod generated mock of class [TestNotifier].
 class MockTestNotifier extends _i1.$Notifier<String>
     with _i2.Mock
     implements _i3.TestNotifier {
   MockTestNotifier(this.seedBuilder);
 
-  late final _$args = ref.$arg as String;
+  late final _$args = ref.$arg as _i5.Model;
 
   String Function() seedBuilder;
 
@@ -26,15 +28,15 @@ class MockTestNotifier extends _i1.$Notifier<String>
   @override
   runBuild() {
     {
-      final ref = this.ref as _i5.$Ref<String, String>;
-      final element = ref.element as _i5.$ClassProviderElement<
-          _i5.AnyNotifier<String, String>, String, Object?, Object?>;
+      final ref = this.ref as _i6.$Ref<String, String>;
+      final element = ref.element as _i6.$ClassProviderElement<
+          _i6.AnyNotifier<String, String>, String, Object?, Object?>;
       element.handleCreate(ref, () => build(_$args));
     }
   }
 
   @override
-  String build(String? param) {
+  String build(_i7.Model? param) {
     return noSuchMethod(Invocation.method(#build, [param]),
         returnValueForMissingStub: seedBuilder());
   }
@@ -44,20 +46,20 @@ class MockTestNotifier extends _i1.$Notifier<String>
     return noSuchMethod(
       Invocation.method(#someMethod, []),
       returnValueForMissingStub:
-          _i6.dummyValue<String>(this, Invocation.method(#someMethod, [])),
+          _i8.dummyValue<String>(this, Invocation.method(#someMethod, [])),
       returnValue:
-          _i6.dummyValue<String>(this, Invocation.method(#someMethod, [])),
+          _i8.dummyValue<String>(this, Invocation.method(#someMethod, [])),
     );
   }
 
   @override
-  _i7.Future<int> someAsyncMethod() async {
+  _i9.Future<int> someAsyncMethod() async {
     return noSuchMethod(
       Invocation.method(#someAsyncMethod, []),
       returnValueForMissingStub:
-          _i6.dummyValue<int>(this, Invocation.method(#someAsyncMethod, [])),
+          _i8.dummyValue<int>(this, Invocation.method(#someAsyncMethod, [])),
       returnValue:
-          _i6.dummyValue<int>(this, Invocation.method(#someAsyncMethod, [])),
+          _i8.dummyValue<int>(this, Invocation.method(#someAsyncMethod, [])),
     );
   }
 
@@ -69,7 +71,7 @@ class MockTestNotifier extends _i1.$Notifier<String>
   }
 
   @override
-  _i7.Future<void> asyncSideEffect() async {
+  _i9.Future<void> asyncSideEffect() async {
     return noSuchMethod(
       Invocation.method(#asyncSideEffect, []),
     );

test/golden/flumepod/family_notifier/family_notifier_test.dart

@@ -1,10 +1,17 @@
+import 'dart:io';
+
 import '../../utils/test_golden_builder.dart';
 
-void main() {
+void main() async {
   testGoldenBuilder(
     "Flumepod - generates standard FamilyNotifier as expected",
     "test/golden/flumepod/family_notifier/input.dart",
     "test/golden/flumepod/family_notifier/input.g.dart",
     "test/golden/flumepod/family_notifier/expected.flumepod.dart",
+    additionalSources: {
+      "flumepod|test/model.dart": await File(
+        "test/golden/flumepod/family_notifier/model.dart",
+      ).readAsString(),
+    },
   );
 }

test/golden/flumepod/family_notifier/input.dart

@@ -1,5 +1,6 @@
 import 'package:flumepod/flumepod.dart';
 import 'package:riverpod_annotation/riverpod_annotation.dart';
+import './model.dart';
 
 part 'input.g.dart';
 
@@ -9,7 +10,7 @@ void main() {}
 @Riverpod(keepAlive: true)
 class TestNotifier extends _$TestNotifier {
   @override
-  String build(String param) {
+  String build(Model param) {
     return "Dummy string";
   }
 

test/golden/flumepod/family_notifier/model.dart

@@ -0,0 +1,5 @@
+class Model {
+  final int param;
+
+  const Model({required this.param});
+}

test/golden/utils/test_golden_builder.dart

@@ -10,8 +10,9 @@ void testGoldenBuilder(
   String description,
   String inputPath,
   String generatedPartPath,
-  String expectedPath,
-) =>
+  String expectedPath, {
+  Map<String, String>? additionalSources,
+}) =>
     test(
       description,
       () async {
@@ -26,6 +27,7 @@ void testGoldenBuilder(
             "flumepod|test/input.dart": inputSource,
             "flumepod|test/input.g.dart": generatedSource,
             "flumepod|test/input.flumepod.dart": expectedSource,
+            ...?additionalSources,
           },
           generateFor: {"flumepod|test/input.dart"},
           outputs: {