From a6e90596b137cfdcb8343669fbe19bf3a61cdc7e Mon Sep 17 00:00:00 2001
From: MitulShah1 <shahmitul005@gmail.com>
Date: Mon, 14 Apr 2025 09:46:18 +0530
Subject: [PATCH] Updated Testcontainer with generic one

---
 aerospike/aerospike_test.go | 86 +++++++++++++++++++++++++++----------
 1 file changed, 64 insertions(+), 22 deletions(-)

diff --git a/aerospike/aerospike_test.go b/aerospike/aerospike_test.go
index 52a525cb..d144b419 100644
--- a/aerospike/aerospike_test.go
+++ b/aerospike/aerospike_test.go
@@ -2,35 +2,73 @@ package aerospike
 
 import (
 	"context"
+	"fmt"
+	"os"
 	"testing"
+	"time"
 
 	"github.com/aerospike/aerospike-client-go/v8"
 	"github.com/stretchr/testify/require"
 	"github.com/testcontainers/testcontainers-go"
-	tcaerospike "github.com/testcontainers/testcontainers-go/modules/aerospike"
+	"github.com/testcontainers/testcontainers-go/wait"
 )
 
 const (
-	aerospikeImage     = "aerospike/aerospike-server:latest"
-	aerospikePort      = "3000/tcp"
-	aerospikeNamespace = "test"
+	// aerospikeImage is the default image used for running Aerospike in tests.
+	aerospikeImage              = "aerospike/aerospike-server:latest"
+	aerospikeImageEnvVar string = "TEST_AEROSPIKE_IMAGE"
+	aerospikePort               = "3000/tcp"
+	febricPort                  = "3001/tcp"
+	heartbeatPort               = "3002/tcp"
+	infoPort                    = "3003/tcp"
+	aerospikeReadyLog           = "migrations: complete"
+	aerospikeNamespace          = "test"
 )
 
-// setupAerospikeTestClient sets up a test Aerospike client
-// using testcontainers-go. It starts an Aerospike container and returns
-// a Storage instance configured to connect to the container.
-// The container is cleaned up after the test completes.
-func setupAerospikeTestClient(t *testing.T) *Storage {
-	t.Helper()
-
-	ctx := context.Background()
-
-	c, err := tcaerospike.Run(ctx, aerospikeImage)
-	if err != nil {
-		t.Fatalf("Failed to start aerospike container: %v", err)
+// startAerospikeContainer starts an Aerospike container for testing
+func startAerospikeContainer(ctx context.Context) (testcontainers.Container, error) {
+	// Get custom image from env if specified
+	image := aerospikeImage
+	if envImage := os.Getenv(aerospikeImageEnvVar); envImage != "" {
+		image = envImage
 	}
 
-	testcontainers.CleanupContainer(t, c)
+	// Container config
+	req := testcontainers.ContainerRequest{
+		Image:        image,
+		ExposedPorts: []string{aerospikePort, febricPort, heartbeatPort, infoPort},
+		WaitingFor: wait.ForAll(
+			wait.ForLog(aerospikeReadyLog),
+			wait.ForListeningPort(aerospikePort).WithStartupTimeout(5*time.Second),
+			wait.ForListeningPort(febricPort).WithStartupTimeout(5*time.Second),
+			wait.ForListeningPort(heartbeatPort).WithStartupTimeout(5*time.Second),
+		),
+		Cmd: []string{
+			"--config-file",
+			"/etc/aerospike/aerospike.conf",
+		},
+	}
+
+	// Start container
+	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
+		ContainerRequest: req,
+		Started:          true,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to start aerospike container: %w", err)
+	}
+
+	return container, nil
+}
+
+// setupTestClient creates a client connected to the test container
+func setupTestClient(t *testing.T) *Storage {
+	t.Helper()
+
+	c, err := startAerospikeContainer(context.Background())
+	if err != nil {
+		t.Fatalf("Failed to start Aerospike container: %v", err)
+	}
 
 	// Extract host and port
 	host, err := c.Host(context.TODO())
@@ -43,12 +81,16 @@ func setupAerospikeTestClient(t *testing.T) *Storage {
 		t.Fatalf("Failed to get container port: %v", err)
 	}
 
+	testcontainers.CleanupContainer(t, c)
+	if err != nil {
+		t.Fatalf("Failed to cleanup Aerospike container: %v", err)
+	}
+
 	return New(Config{
 		Hosts:     []*aerospike.Host{aerospike.NewHost(host, port.Int())},
 		Reset:     true,
 		Namespace: aerospikeNamespace,
 	})
-
 }
 
 // Test_AeroSpikeDB_Get tests the Get method
@@ -58,7 +100,7 @@ func Test_AeroSpikeDB_Get(t *testing.T) {
 		val = []byte("doe")
 	)
 
-	testStore := setupAerospikeTestClient(t)
+	testStore := setupTestClient(t)
 	defer testStore.Close()
 
 	// Set a value
@@ -79,7 +121,7 @@ func Test_AeroSpikeDB_Delete(t *testing.T) {
 		val = []byte("doe")
 	)
 
-	testStore := setupAerospikeTestClient(t)
+	testStore := setupTestClient(t)
 	defer testStore.Close()
 
 	// Set a value
@@ -106,7 +148,7 @@ func Test_AeroSpikeDB_Reset(t *testing.T) {
 		val2 = []byte("smith")
 	)
 
-	testStore := setupAerospikeTestClient(t)
+	testStore := setupTestClient(t)
 	defer testStore.Close()
 
 	// Set multiple values
@@ -132,7 +174,7 @@ func Test_AeroSpikeDB_Reset(t *testing.T) {
 // Test_AeroSpikeDB_GetSchemaInfo tests the GetSchemaInfo method
 func Test_AeroSpikeDB_GetSchemaInfo(t *testing.T) {
 
-	testStore := setupAerospikeTestClient(t)
+	testStore := setupTestClient(t)
 	defer testStore.Close()
 
 	// Get schema info