wehub-resource-sync
100 lines
2.7 KiB
Go
100 lines
2.7 KiB
Go
package distributedhdr_test
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"testing"
|
|
|
|
"github.com/mudler/LocalAI/pkg/distributedhdr"
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
)
|
|
|
|
func TestDistributedHdr(t *testing.T) {
|
|
RegisterFailHandler(Fail)
|
|
RunSpecs(t, "distributedhdr suite")
|
|
}
|
|
|
|
var _ = Describe("distributedhdr", func() {
|
|
It("returns empty string when no holder is attached", func() {
|
|
Expect(distributedhdr.Holder(context.Background())).To(BeNil())
|
|
Expect(distributedhdr.Load(distributedhdr.Holder(context.Background()))).To(BeEmpty())
|
|
})
|
|
|
|
It("is a no-op when stamping into a context without a holder", func() {
|
|
Expect(func() {
|
|
distributedhdr.Stamp(context.Background(), "worker-x")
|
|
}).ToNot(Panic())
|
|
})
|
|
|
|
It("is a no-op when stamping a nil context", func() {
|
|
Expect(func() {
|
|
distributedhdr.Stamp(nil, "worker-x")
|
|
}).ToNot(Panic())
|
|
})
|
|
|
|
It("stores and retrieves a node ID through Stamp/Load", func() {
|
|
h := distributedhdr.NewHolder()
|
|
ctx := distributedhdr.WithHolder(context.Background(), h)
|
|
|
|
distributedhdr.Stamp(ctx, "worker-a")
|
|
Expect(distributedhdr.Load(h)).To(Equal("worker-a"))
|
|
})
|
|
|
|
It("retains the latest stamp when called twice", func() {
|
|
h := distributedhdr.NewHolder()
|
|
ctx := distributedhdr.WithHolder(context.Background(), h)
|
|
|
|
distributedhdr.Stamp(ctx, "worker-a")
|
|
distributedhdr.Stamp(ctx, "worker-b")
|
|
Expect(distributedhdr.Load(h)).To(Equal("worker-b"))
|
|
})
|
|
|
|
It("does not overwrite when stamping with an empty string", func() {
|
|
h := distributedhdr.NewHolder()
|
|
ctx := distributedhdr.WithHolder(context.Background(), h)
|
|
|
|
distributedhdr.Stamp(ctx, "worker-a")
|
|
distributedhdr.Stamp(ctx, "")
|
|
Expect(distributedhdr.Load(h)).To(Equal("worker-a"))
|
|
})
|
|
|
|
It("propagates the holder via Inherit into a derived context", func() {
|
|
h := distributedhdr.NewHolder()
|
|
src := distributedhdr.WithHolder(context.Background(), h)
|
|
dst := distributedhdr.Inherit(context.Background(), src)
|
|
|
|
distributedhdr.Stamp(dst, "worker-c")
|
|
Expect(distributedhdr.Load(h)).To(Equal("worker-c"))
|
|
})
|
|
|
|
It("Inherit is a no-op when src has no holder", func() {
|
|
dst := distributedhdr.Inherit(context.Background(), context.Background())
|
|
Expect(distributedhdr.Holder(dst)).To(BeNil())
|
|
})
|
|
|
|
It("publishes stamps across goroutines race-free", func() {
|
|
h := distributedhdr.NewHolder()
|
|
ctx := distributedhdr.WithHolder(context.Background(), h)
|
|
|
|
const N = 64
|
|
var wg sync.WaitGroup
|
|
wg.Add(N)
|
|
for i := 0; i < N; i++ {
|
|
i := i
|
|
go func() {
|
|
defer wg.Done()
|
|
if i%2 == 0 {
|
|
distributedhdr.Stamp(ctx, "worker-even")
|
|
} else {
|
|
distributedhdr.Stamp(ctx, "worker-odd")
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
|
|
got := distributedhdr.Load(h)
|
|
Expect(got).To(SatisfyAny(Equal("worker-even"), Equal("worker-odd")))
|
|
})
|
|
})
|