From 4e0f89aed60510dbb3926aff9623ea1cdf2d4de2 Mon Sep 17 00:00:00 2001
From: Fedorov Alexander <alfedorov@ozon.ru>
Date: Mon, 28 Nov 2022 10:45:23 +0300
Subject: [PATCH] [ISSUE-45] fix waiting for async steps in parent step

---
 pkg/framework/core/common/step_context.go | 2 +-
 pkg/framework/core/common/steps.go        | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/pkg/framework/core/common/step_context.go b/pkg/framework/core/common/step_context.go
index 49b4d2e..2db327d 100644
--- a/pkg/framework/core/common/step_context.go
+++ b/pkg/framework/core/common/step_context.go
@@ -150,6 +150,7 @@ func (ctx *stepCtx) WithNewStep(stepName string, step func(ctx provider.StepCtx)
 	defer ctx.currentStep.WithChild(newCtx.CurrentStep())
 	defer func() {
 		r := recover()
+		newCtx.WG().Wait()
 		newCtx.CurrentStep().Finish()
 		if r != nil {
 			ctxName := newCtx.ExecutionContextName()
@@ -168,7 +169,6 @@ func (ctx *stepCtx) WithNewAsyncStep(stepName string, step func(ctx provider.Ste
 		wg = ctx.parentStep.WG()
 	}
 	wg.Add(1)
-	defer wg.Wait()
 
 	go func() {
 		defer wg.Done()
diff --git a/pkg/framework/core/common/steps.go b/pkg/framework/core/common/steps.go
index 13efbcb..9d578fb 100644
--- a/pkg/framework/core/common/steps.go
+++ b/pkg/framework/core/common/steps.go
@@ -15,6 +15,7 @@ func (c *Common) WithNewStep(stepName string, step func(ctx provider.StepCtx), p
 	defer c.Step(stCtx.CurrentStep())
 	defer func() {
 		r := recover()
+		stCtx.WG().Wait()
 		stCtx.CurrentStep().Finish()
 		if r != nil {
 			ctxName := c.ExecutionContext().GetName()