From e055090b6870bed143e2e74e5c89342708f2b0ed Mon Sep 17 00:00:00 2001
From: Sanyam Singhal <ssinghal@yugabyte.com>
Date: Wed, 5 Mar 2025 14:01:04 +0000
Subject: [PATCH] review comments

---
 yb-voyager/cmd/assessMigrationCommand_test.go | 68 +++++++++++++++++--
 yb-voyager/test/utils/cmdutils.go             |  2 +
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/yb-voyager/cmd/assessMigrationCommand_test.go b/yb-voyager/cmd/assessMigrationCommand_test.go
index 4b83ea5a1..e824d46d9 100644
--- a/yb-voyager/cmd/assessMigrationCommand_test.go
+++ b/yb-voyager/cmd/assessMigrationCommand_test.go
@@ -27,6 +27,8 @@ import (
 
 	"github.com/davecgh/go-spew/spew"
 	log "github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/yugabyte/yb-voyager/yb-voyager/src/migassessment"
 	"github.com/yugabyte/yb-voyager/yb-voyager/src/utils"
 	testcontainers "github.com/yugabyte/yb-voyager/yb-voyager/test/containers"
 	testutils "github.com/yugabyte/yb-voyager/yb-voyager/test/utils"
@@ -38,6 +40,10 @@ func TestMain(m *testing.M) {
 	log.SetLevel(log.WarnLevel)
 
 	exitCode := m.Run()
+
+	// cleanig up all the running containers
+	testcontainers.TerminateAllContainers()
+
 	os.Exit(exitCode)
 }
 
@@ -52,7 +58,6 @@ func Test_AssessMigration(t *testing.T) {
 	if err != nil {
 		utils.ErrExit("Failed to start postgres container: %v", err)
 	}
-	defer postgresContainer.Terminate(context.Background())
 
 	// create table and initial data in it
 	postgresContainer.ExecuteSqls(`
@@ -82,6 +87,49 @@ END $$;`)
 		}
 	}
 
+	expectedColocatedTables := []string{"public.test_data"}
+	expectedSizingAssessmentReport := migassessment.SizingAssessmentReport{
+		SizingRecommendation: migassessment.SizingRecommendation{
+			ColocatedTables:                 []string{"public.test_data"},
+			ColocatedReasoning:              "Recommended instance type with 4 vCPU and 16 GiB memory could fit 1 objects (1 tables/materialized views and 0 explicit/implicit indexes) with 6.52 MB size and throughput requirement of 10 reads/sec and 0 writes/sec as colocated. Non leaf partition tables/indexes and unsupported tables/indexes were not considered.",
+			ShardedTables:                   nil,
+			NumNodes:                        3,
+			VCPUsPerInstance:                4,
+			MemoryPerInstance:               16,
+			OptimalSelectConnectionsPerNode: 8,
+			OptimalInsertConnectionsPerNode: 12,
+			EstimatedTimeInMinForImport:     1,
+			ParallelVoyagerJobs:             1,
+		},
+		FailureReasoning: "",
+	}
+	expectedTableIndexStats := []migassessment.TableIndexStats{
+		{
+			SchemaName:      "public",
+			ObjectName:      "test_data",
+			RowCount:        int64Ptr(100000),
+			ColumnCount:     int64Ptr(2),
+			ReadsPerSecond:  int64Ptr(10),
+			WritesPerSecond: int64Ptr(0),
+			IsIndex:         false,
+			ObjectType:      "table",
+			ParentTableName: nil,
+			SizeInBytes:     int64Ptr(6832128),
+		},
+		{
+			SchemaName:      "public",
+			ObjectName:      "test_data_pkey",
+			RowCount:        nil,
+			ColumnCount:     int64Ptr(1),
+			ReadsPerSecond:  int64Ptr(0),
+			WritesPerSecond: int64Ptr(0),
+			IsIndex:         true,
+			ObjectType:      "primary key",
+			ParentTableName: stringPtr("public.test_data"),
+			SizeInBytes:     int64Ptr(2260992),
+		},
+	}
+
 	// running the command
 	err = testutils.RunVoyagerCommmand(postgresContainer, "assess-migration", []string{
 		"--iops-capture-interval", "10",
@@ -103,8 +151,20 @@ END $$;`)
 		t.Errorf("failed to get colocated tables recommendation: %v", err)
 	}
 
-	fmt.Printf("Colocated tables: %v\n", colocatedTables)
+	log.Infof("Colocated tables: %v\n", colocatedTables)
+	log.Infof("Sizing recommendations: %+v\n", report.Sizing)
+	log.Infof("TableIndexStats: %s\n", spew.Sdump(report.TableIndexStats))
+
+	// assert all these three info
+	testutils.AssertEqualStringSlices(t, expectedColocatedTables, colocatedTables)
+	assert.Equal(t, expectedSizingAssessmentReport, *report.Sizing)
+	assert.Equal(t, expectedTableIndexStats, *report.TableIndexStats)
+}
+
+func int64Ptr(i int64) *int64 {
+	return &i
+}
 
-	fmt.Printf("Sizing recommendations: %+v\n", report.Sizing)
-	fmt.Printf("TableIndexStats: %s\n", spew.Sdump(report.TableIndexStats))
+func stringPtr(s string) *string {
+	return &s
 }
diff --git a/yb-voyager/test/utils/cmdutils.go b/yb-voyager/test/utils/cmdutils.go
index 928e75008..343c8e54e 100644
--- a/yb-voyager/test/utils/cmdutils.go
+++ b/yb-voyager/test/utils/cmdutils.go
@@ -73,6 +73,8 @@ func RunVoyagerCommmand(container testcontainers.TestContainer,
 	cmdArgs = append(connectionArgs, cmdArgs...)
 	cmdStr := fmt.Sprintf("yb-voyager %s %s", cmdName, strings.Join(cmdArgs, " "))
 	cmd := exec.Command("/bin/bash", "-c", cmdStr)
+	// don't send to callhome for tests
+	cmd.Env = append(os.Environ(), "YB_VOYAGER_SEND_DIAGNOSTICS=false")
 
 	// 2) Get the stdout and stderr pipes.
 	stdoutPipe, err := cmd.StdoutPipe()