dify-plugin-daemon/internal/cluster/vote_test.go

package cluster

import (
	"fmt"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
)

func createSimulationHealthCheckSever() (uint16, error) {
	gin.SetMode(gin.ReleaseMode)
	router := gin.Default()

	router.GET("/health/check", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "ok"})
	})

	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		return 0, fmt.Errorf("failed to get a random port: %s", err.Error())
	}
	listener.Close()

	port := listener.Addr().(*net.TCPAddr).Port

	go func() {
		router.Run(fmt.Sprintf(":%d", port))
	}()

	return uint16(port), nil
}

func TestVoteIps(t *testing.T) {
	// create a health check server
	port, err := createSimulationHealthCheckSever()
	if err != nil {
		t.Errorf("create simulation health check server failed: %v", err)
		return
	}

	cluster, err := createSimulationCluster(2)
	if err != nil {
		t.Errorf("create simulation cluster failed: %v", err)
		return
	}

	for _, node := range cluster {
		node.port = port
	}

	wg := sync.WaitGroup{}
	wg.Add(2)

	// wait for all voting processes complete
	for _, node := range cluster {
		node := node
		go func() {
			defer wg.Done()
			<-node.NotifyVotingCompleted()
		}()
	}

	launchSimulationCluster(cluster, t)
	defer closeSimulationCluster(cluster, t)

	// wait for all nodes to be ready
	wg.Wait()

	// wait for all ips to be voted
	time.Sleep(time.Second)

	for _, node := range cluster {
		nodes, err := node.GetNodes()
		if err != nil {
			t.Errorf("get nodes failed: %v", err)
			return
		}

		for _, node := range nodes {
			for _, ip := range node.Ips {
				if len(ip.Votes) == 0 {
					t.Errorf("vote for ip %s failed", ip.Address)
					return
				}
			}
		}
	}
}