fix tests

base/commands/migration/cancel_it_test.go

@@ -7,13 +7,13 @@ import (
 	"os"
 	"sync"
 	"testing"
-	"time"
 
 	_ "github.com/hazelcast/hazelcast-commandline-client/base"
 	_ "github.com/hazelcast/hazelcast-commandline-client/base/commands"
 	"github.com/hazelcast/hazelcast-commandline-client/base/commands/migration"
 	. "github.com/hazelcast/hazelcast-commandline-client/internal/check"
 	"github.com/hazelcast/hazelcast-commandline-client/internal/it"
+	hz "github.com/hazelcast/hazelcast-go-client"
 	"github.com/hazelcast/hazelcast-go-client/serialization"
 	"github.com/stretchr/testify/require"
 )
@@ -45,16 +45,18 @@ func cancelTest(t *testing.T) {
 	tcx := it.TestContext{T: t}
 	ctx := context.Background()
 	tcx.Tester(func(tcx it.TestContext) {
+		ci := hz.NewClientInternal(tcx.Client)
 		mID := migrationIDFunc()
 		createMapping(ctx, tcx)
 		setStatusInProgress(tcx, ctx)
 		var wg sync.WaitGroup
 		wg.Add(1)
-		go tcx.WithReset(func() {
+		go func() {
 			defer wg.Done()
 			Must(tcx.CLC().Execute(ctx, "cancel"))
-		})
-		time.Sleep(1 * time.Second)
+		}()
+		cq := MustValue(ci.Client().GetQueue(ctx, migration.CancelQueue))
+		MustValue(cq.Poll(ctx))
 		setStatusCancelling(mID, tcx, ctx)
 		wg.Wait()
 		statusMap := MustValue(tcx.Client.GetMap(ctx, migration.StatusMapName))

base/commands/migration/status_stages_it_test.go

@@ -8,7 +8,6 @@ import (
 	"os"
 	"sync"
 	"testing"
-	"time"
 
 	_ "github.com/hazelcast/hazelcast-commandline-client/base"
 	_ "github.com/hazelcast/hazelcast-commandline-client/base/commands"
@@ -63,14 +62,13 @@ func statusTest(t *testing.T) {
 		mID := setStatusInProgress(tcx, ctx)
 		var wg sync.WaitGroup
 		wg.Add(1)
-		go tcx.WithReset(func() {
+		go func() {
 			defer wg.Done()
 			Must(tcx.CLC().Execute(ctx, "status", "-o", outDir))
-		})
-		time.Sleep(1 * time.Second)
+		}()
+		tcx.AssertStdoutContains("Connected to the migration cluster")
 		setStatusCompleted(mID, tcx, ctx)
 		wg.Wait()
-		tcx.AssertStdoutContains("Connected to the migration cluster")
 		tcx.WithReset(func() {
 			f := paths.Join(outDir, fmt.Sprintf("migration_report_%s.txt", mID))
 			require.Equal(t, true, paths.Exists(f))