From 2ce8d778360d8c0d11c28a47e475f4a223296cfb Mon Sep 17 00:00:00 2001
From: Jacob MacDonald <jakemac@google.com>
Date: Tue, 22 Oct 2024 09:06:31 -0700
Subject: [PATCH] fix addModel and add test, expose current model (#110)

Follow up to https://github.com/dart-lang/macros/pull/105#discussion_r1810154444
---
 pkgs/dart_model/lib/src/scopes.dart   |  7 ++++++-
 pkgs/dart_model/test/scopes_test.dart | 20 +++++++++++++++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/pkgs/dart_model/lib/src/scopes.dart b/pkgs/dart_model/lib/src/scopes.dart
index c2cb609f..e2094c7e 100644
--- a/pkgs/dart_model/lib/src/scopes.dart
+++ b/pkgs/dart_model/lib/src/scopes.dart
@@ -159,12 +159,17 @@ class MacroScope {
   /// multiple queries.
   void addModel(Model model) {
     if (_accumulatedModel case var accumulated?) {
-      accumulated.mergeWith(model);
+      _accumulatedModel = accumulated.mergeWith(model);
     } else {
       _accumulatedModel = model;
     }
   }
 
+  /// The current accumulated model for this macro scope.
+  ///
+  /// This is only safe to use after [addModel] has been called at least once.
+  Model get model => _accumulatedModel!;
+
   static MacroScope get current {
     final scope = Scope._currentOrNull;
     return switch (scope) {
diff --git a/pkgs/dart_model/test/scopes_test.dart b/pkgs/dart_model/test/scopes_test.dart
index 98cce985..0e112de0 100644
--- a/pkgs/dart_model/test/scopes_test.dart
+++ b/pkgs/dart_model/test/scopes_test.dart
@@ -7,7 +7,7 @@ import 'package:dart_model/src/json_buffer/json_buffer_builder.dart';
 import 'package:test/test.dart';
 
 void main() {
-  group(Scope, () {
+  group('Scope', () {
     for (final scope in [Scope.none, Scope.macro, Scope.query]) {
       test('create maps and serialize work in $scope', () {
         expect(scope.run(() => Scope.createMap(TypedMapSchema({}))),
@@ -33,5 +33,23 @@ void main() {
       Scope.macro.run(() => Scope.none.run(() {}));
       Scope.query.run(() => Scope.none.run(() {}));
     });
+
+    test('addModel merges the model with the previous', () {
+      late Model initial;
+      late Model firstQuery;
+      Scope.query.run(() {
+        initial = Model()..uris['a'] = (Library()..scopes['A'] = Interface());
+        firstQuery = Model()
+          ..uris['a'] = (Library()..scopes['B'] = Interface());
+      });
+      Scope.macro.run(() {
+        MacroScope.current.addModel(initial);
+        expect(MacroScope.current.model.uris['a']!.scopes['A'], isNot(null));
+        expect(MacroScope.current.model.uris['a']!.scopes['B'], null);
+        MacroScope.current.addModel(firstQuery);
+        expect(MacroScope.current.model.uris['a']!.scopes['A'], isNot(null));
+        expect(MacroScope.current.model.uris['a']!.scopes['B'], isNot(null));
+      });
+    });
   });
 }