diff --git a/packages/mst-query/tests/models/RootStore.ts b/packages/mst-query/tests/models/RootStore.ts
index 10d3547..6be3691 100644
--- a/packages/mst-query/tests/models/RootStore.ts
+++ b/packages/mst-query/tests/models/RootStore.ts
@@ -14,6 +14,7 @@ import { ArrayQuery } from './ArrayQuery';
 import { SafeReferenceQuery } from './SafeReferenceQuery';
 import { RemoveItemMutation } from './RemoveItemMutation';
 import { ErrorMutation } from './ErrorMutation';
+import { FixedModel, FormatModel } from './UnionModel';
 
 export const DateModel = types.model('DateModel', {
     id: types.identifier,
@@ -48,8 +49,8 @@ const AmountLimitModel = types.model('AmountLimit').props({
             types.model({
                 tag: types.enumeration(Object.values(AmountTag)),
                 content: types.maybeNull(types.string),
-            })
-        )
+            }),
+        ),
     ),
 });
 
@@ -59,7 +60,7 @@ const TestModel = types.model({
     prop: types.maybeNull(
         types.model({
             ids: types.array(types.model({ baha: types.string })),
-        })
+        }),
     ),
     folderPath: types.maybe(types.string),
     origin: types.union(types.string, types.undefined),
@@ -105,6 +106,8 @@ export const Root = createRootStore({
     listStore: optional(createModelStore('ListStore', ListModel)),
     dateStore: optional(createModelStore('DateStore', DateModel)),
     deepModelCStore: optional(createModelStore('DeepModelCStore', DeepModelC)),
+    fixedModelStore: optional(createModelStore('FixedModelStore', FixedModel)),
+    formatModelStore: optional(createModelStore('FixedModelStore', FormatModel)),
     serviceStore: optional(ServiceStore),
 });
 
diff --git a/packages/mst-query/tests/models/UnionModel.ts b/packages/mst-query/tests/models/UnionModel.ts
new file mode 100644
index 0000000..7dc5ef4
--- /dev/null
+++ b/packages/mst-query/tests/models/UnionModel.ts
@@ -0,0 +1,15 @@
+import { types } from 'mobx-state-tree';
+
+export const FixedModel = types.model('FixedModel', {
+    id: types.identifier,
+    kind: types.literal('FIXED'),
+    fixedValue: types.string,
+});
+
+export const FormatModel = types.model('FormatModel', {
+    id: types.identifier,
+    kind: types.literal('FORMAT'),
+    formatValue: types.string,
+});
+
+export const UnionModel = types.union(FixedModel, FormatModel);
diff --git a/packages/mst-query/tests/mstQuery.test.tsx b/packages/mst-query/tests/mstQuery.test.tsx
index 58c2246..85c087e 100644
--- a/packages/mst-query/tests/mstQuery.test.tsx
+++ b/packages/mst-query/tests/mstQuery.test.tsx
@@ -15,6 +15,7 @@ import { QueryClient } from '../src/QueryClient';
 import { createContext } from '../src/QueryClientProvider';
 import { DateModel, DeepModelA, Root } from './models/RootStore';
 import { useInfiniteQuery, useVolatileQuery } from '../src/hooks';
+import { UnionModel } from './models/UnionModel';
 
 const setup = ({ strictMode = false } = {}) => {
     const queryClient = new QueryClient({ RootStore: Root });
@@ -1309,33 +1310,26 @@ test('initial data should only be set on mount', async () => {
 });
 
 test('union of array models', () => {
+    const { queryClient } = setup();
+
     const data = {
         rules: [
             {
+                id: '1',
                 kind: 'FIXED',
                 fixedValue: 'Fixed value',
             },
             {
+                id: '2',
                 kind: 'FORMAT',
                 formatValue: 'Formatted value',
-            },            
+            },
         ],
     };
     const Model = types.model('UnionArrayTestModel', {
-        rules: types.array(
-            types.union(
-                types.late(() => types.model('FixedModel', {
-                    kind: types.literal('FIXED'),
-                    fixedValue: types.string,
-                })),
-                types.model('FormatModel', {
-                    kind: types.literal('FORMAT'),
-                    formatValue: types.string,
-                }),
-            ),
-        ),
+        rules: types.array(types.reference(UnionModel)),
     });
-    const result = merge(data, Model, {});
+    const result = merge(data, Model, queryClient.config.env);
     expect(result.rules[0].fixedValue).toBe('Fixed value');
     expect(result.rules[1].formatValue).toBe('Formatted value');
 });