From f7d6c7fe357030474dfb01ec55799526bfa56584 Mon Sep 17 00:00:00 2001
From: pgayvallet <pierre.gayvallet@elastic.co>
Date: Mon, 27 Feb 2023 11:18:49 +0100
Subject: [PATCH] improve `ModelStage` utility type

---
 .../src/zdt/model/stages/init.ts                      |  2 +-
 .../src/zdt/model/types.ts                            | 11 ++++++++---
 .../src/zdt/state/index.ts                            |  1 +
 .../src/zdt/state/types.ts                            |  7 ++++++-
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/model/stages/init.ts b/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/model/stages/init.ts
index 48aa0824f6411..78dccf237afca 100644
--- a/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/model/stages/init.ts
+++ b/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/model/stages/init.ts
@@ -13,7 +13,7 @@ import { isTypeof } from '../../actions';
 import type { State } from '../../state';
 import type { ModelStage } from '../types';
 
-export const init: ModelStage<'INIT'> = (state, res, context): State => {
+export const init: ModelStage<'INIT', 'DONE' | 'FATAL'> = (state, res, context): State => {
   if (Either.isLeft(res)) {
     const left = res.left;
     if (isTypeof(left, 'incompatible_cluster_routing_allocation')) {
diff --git a/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/model/types.ts b/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/model/types.ts
index db8e67e183232..6c8227794bb83 100644
--- a/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/model/types.ts
+++ b/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/model/types.ts
@@ -8,7 +8,12 @@
 
 import { ExcludeRetryableEsError } from '../../model/types';
 import type { MigratorContext } from '../context';
-import type { State, AllActionStates, StateFromActionState } from '../state';
+import type {
+  AllActionStates,
+  AllControlStates,
+  StateFromActionState,
+  StateFromControlState,
+} from '../state';
 import type { ResponseType } from '../next';
 
 /**
@@ -21,8 +26,8 @@ export type StateActionResponse<T extends AllActionStates> = ExcludeRetryableEsE
 /**
  * Defines a stage delegation function for the model
  */
-export type ModelStage<T extends AllActionStates> = (
+export type ModelStage<T extends AllActionStates, R extends AllControlStates> = (
   state: StateFromActionState<T>,
   res: StateActionResponse<T>,
   context: MigratorContext
-) => State;
+) => StateFromControlState<T | R>;
diff --git a/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/state/index.ts b/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/state/index.ts
index 2954d7f2f9dd2..bff3ea15da4c2 100644
--- a/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/state/index.ts
+++ b/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/state/index.ts
@@ -15,5 +15,6 @@ export type {
   AllActionStates,
   AllControlStates,
   StateFromActionState,
+  StateFromControlState,
 } from './types';
 export { createInitialState } from './create_initial_state';
diff --git a/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/state/types.ts b/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/state/types.ts
index 37a93de1eee41..90a888a8a947c 100644
--- a/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/state/types.ts
+++ b/packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/state/types.ts
@@ -49,4 +49,9 @@ export interface ControlStateMap {
 /**
  * Utility type to reverse lookup an `AllControlStates` to it's corresponding State subtype.
  */
-export type StateFromActionState<T extends AllActionStates> = ControlStateMap[T];
+export type StateFromControlState<T extends AllControlStates> = ControlStateMap[T];
+
+/**
+ * Utility type to reverse lookup an `AllActionStates` to it's corresponding State subtype.
+ */
+export type StateFromActionState<T extends AllActionStates> = StateFromControlState<T>;