wehub-resource-sync
38 lines
1.4 KiB
Go
38 lines
1.4 KiB
Go
package distributedhdr_test
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/mudler/LocalAI/pkg/distributedhdr"
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
)
|
|
|
|
var _ = Describe("prefix chain ctx", func() {
|
|
It("round-trips the chain through ctx", func() {
|
|
ctx := distributedhdr.WithPrefixChain(context.Background(), []uint64{1, 2, 3})
|
|
Expect(distributedhdr.PrefixChain(ctx)).To(Equal([]uint64{1, 2, 3}))
|
|
})
|
|
It("returns nil when absent", func() {
|
|
Expect(distributedhdr.PrefixChain(context.Background())).To(BeNil())
|
|
})
|
|
|
|
It("uses the hook to build the chain when set", func() {
|
|
distributedhdr.PrefixChainHook = func(model, prompt string) []uint64 { return []uint64{42} }
|
|
defer func() { distributedhdr.PrefixChainHook = nil }()
|
|
ctx := distributedhdr.MaybeWithPrefixChain(context.Background(), "m", "hi")
|
|
Expect(distributedhdr.PrefixChain(ctx)).To(Equal([]uint64{42}))
|
|
})
|
|
It("is a no-op when the hook is nil", func() {
|
|
distributedhdr.PrefixChainHook = nil
|
|
ctx := distributedhdr.MaybeWithPrefixChain(context.Background(), "m", "hi")
|
|
Expect(distributedhdr.PrefixChain(ctx)).To(BeNil())
|
|
})
|
|
It("is a no-op when the hook returns an empty chain", func() {
|
|
distributedhdr.PrefixChainHook = func(model, prompt string) []uint64 { return nil }
|
|
defer func() { distributedhdr.PrefixChainHook = nil }()
|
|
ctx := distributedhdr.MaybeWithPrefixChain(context.Background(), "m", "hi")
|
|
Expect(distributedhdr.PrefixChain(ctx)).To(BeNil())
|
|
})
|
|
})
|