LocalAI/pkg/model/store_test.go

package model_test

import (
	"github.com/mudler/LocalAI/pkg/model"
	"github.com/mudler/LocalAI/pkg/system"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

var _ = Describe("InMemoryModelStore", func() {
	var store *model.InMemoryModelStore

	BeforeEach(func() {
		store = model.NewInMemoryModelStore()
	})

	It("Get returns false for missing key", func() {
		m, ok := store.Get("nonexistent")
		Expect(ok).To(BeFalse())
		Expect(m).To(BeNil())
	})

	It("Set then Get returns the model", func() {
		expected := model.NewModel("test-model", "localhost:5000", nil)
		store.Set("test-model", expected)

		m, ok := store.Get("test-model")
		Expect(ok).To(BeTrue())
		Expect(m).To(Equal(expected))
	})

	It("Set overwrites existing entry", func() {
		first := model.NewModel("m1", "addr1", nil)
		second := model.NewModel("m1", "addr2", nil)

		store.Set("m1", first)
		store.Set("m1", second)

		m, ok := store.Get("m1")
		Expect(ok).To(BeTrue())
		Expect(m).To(Equal(second))
	})

	It("Delete removes model, Get returns false", func() {
		store.Set("to-delete", model.NewModel("to-delete", "addr", nil))
		store.Delete("to-delete")

		m, ok := store.Get("to-delete")
		Expect(ok).To(BeFalse())
		Expect(m).To(BeNil())
	})

	It("Delete on missing key does not panic", func() {
		Expect(func() { store.Delete("no-such-key") }).ToNot(Panic())
	})

	It("Range visits all entries", func() {
		store.Set("a", model.NewModel("a", "addr-a", nil))
		store.Set("b", model.NewModel("b", "addr-b", nil))
		store.Set("c", model.NewModel("c", "addr-c", nil))

		visited := map[string]bool{}
		store.Range(func(id string, m *model.Model) bool {
			visited[id] = true
			return true
		})
		Expect(visited).To(HaveLen(3))
		Expect(visited).To(HaveKey("a"))
		Expect(visited).To(HaveKey("b"))
		Expect(visited).To(HaveKey("c"))
	})

	It("Range stops early when callback returns false", func() {
		store.Set("x", model.NewModel("x", "addr-x", nil))
		store.Set("y", model.NewModel("y", "addr-y", nil))
		store.Set("z", model.NewModel("z", "addr-z", nil))

		count := 0
		store.Range(func(id string, m *model.Model) bool {
			count++
			return false // stop after first
		})
		Expect(count).To(Equal(1))
	})

	It("Range on empty store is a no-op", func() {
		called := false
		store.Range(func(id string, m *model.Model) bool {
			called = true
			return true
		})
		Expect(called).To(BeFalse())
	})
})

var _ = Describe("ModelLoader with custom ModelStore", func() {
	var (
		modelLoader *model.ModelLoader
		modelPath   string
		customStore *model.InMemoryModelStore
	)

	BeforeEach(func() {
		modelPath = GinkgoT().TempDir()

		systemState, err := system.GetSystemState(
			system.WithModelPath(modelPath),
		)
		Expect(err).ToNot(HaveOccurred())
		modelLoader = model.NewModelLoader(systemState)

		customStore = model.NewInMemoryModelStore()
	})

	Context("SetModelStore", func() {
		It("ListLoadedModels uses the custom store", func() {
			// Pre-populate the custom store
			customStore.Set("remote-model-1", model.NewModel("remote-model-1", "node1:5000", nil))
			customStore.Set("remote-model-2", model.NewModel("remote-model-2", "node2:5000", nil))

			modelLoader.SetModelStore(customStore)

			listed := modelLoader.ListLoadedModels()
			Expect(listed).To(HaveLen(2))

			ids := map[string]bool{}
			for _, m := range listed {
				ids[m.ID] = true
			}
			Expect(ids).To(HaveKey("remote-model-1"))
			Expect(ids).To(HaveKey("remote-model-2"))
		})

		It("ShutdownModel uses the custom store", func() {
			customStore.Set("shutdown-me", model.NewModel("shutdown-me", "addr", nil))
			modelLoader.SetModelStore(customStore)

			err := modelLoader.ShutdownModel("shutdown-me")
			Expect(err).ToNot(HaveOccurred())

			// Model should be removed from the custom store
			_, ok := customStore.Get("shutdown-me")
			Expect(ok).To(BeFalse())

			// ListLoadedModels should no longer include it
			Expect(modelLoader.ListLoadedModels()).To(BeEmpty())
		})
	})
})