From 50f81b7e387667b9a766d44d5ed9f2a82fb145ce Mon Sep 17 00:00:00 2001
From: Thomas Farr <tsfarr@amazon.com>
Date: Thu, 27 Feb 2025 16:54:06 +1300
Subject: [PATCH] Fix flow framework test

Signed-off-by: Thomas Farr <tsfarr@amazon.com>
---
 spec/schemas/flow_framework.common.yaml       | 27 ++++++++++++-------
 .../flow_framework/workflow/state/search.yaml |  9 ++++---
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/spec/schemas/flow_framework.common.yaml b/spec/schemas/flow_framework.common.yaml
index a478e0261..b6da56940 100644
--- a/spec/schemas/flow_framework.common.yaml
+++ b/spec/schemas/flow_framework.common.yaml
@@ -129,14 +129,6 @@ components:
         skipped:
           type: integer
           description: The number of shards where the deletion was skipped.
-    query:
-      type: object
-      description: The search query to match workflows. Use `match_all` to retrieve all workflows, or `match` to search by specific fields like `use_case`.
-      properties:
-        match:
-          type: object
-        match_all:
-          type: object
     total:
       type: object
       properties:
@@ -205,6 +197,23 @@ components:
           type: string
         user:
           $ref: '#/components/schemas/user'
+        provision_start_time:
+          $ref: '_common.yaml#/components/schemas/DateTime'
+        provision_end_time:
+          $ref: '_common.yaml#/components/schemas/DateTime'
+        resources_created:
+          $ref: '#/components/schemas/ResourcesCreated'
+    ResourcesCreated:
+      type: object
+      properties:
+        workflow_step_name:
+          type: string
+        workflow_step_id:
+          type: string
+        resource_type:
+          type: string
+        resource_id:
+          type: string
     user:
       type: object
       properties:
@@ -348,7 +357,7 @@ components:
       type: object
       properties:
         query:
-          $ref: '#/components/schemas/query'
+          $ref: '_common.query_dsl.yaml#/components/schemas/QueryContainer'
     WorkflowSearchResponse:
       type: object
       properties:
diff --git a/tests/default/flow_framework/workflow/state/search.yaml b/tests/default/flow_framework/workflow/state/search.yaml
index 67138eb1e..5f5aa0f76 100644
--- a/tests/default/flow_framework/workflow/state/search.yaml
+++ b/tests/default/flow_framework/workflow/state/search.yaml
@@ -18,13 +18,15 @@ prologues:
     output:
       workflow_id: payload.workflow_id
 chapters:
-  - synopsis: Search for a workflow.
+  - synopsis: Search for a workflow state.
     path: /_plugins/_flow_framework/workflow/state/_search
     method: GET
     request:
       payload:
         query:
-          match_all: {}
+          term:
+            workflow_id:
+              value: ${workflow.workflow_id}
     response:
       status: 200
       payload:
@@ -32,5 +34,4 @@ chapters:
           hits:
             - _index: .plugins-flow-framework-state
               _source:
-                state: NOT_STARTED
-            
+                workflow_id: ${workflow.workflow_id}
\ No newline at end of file