routes_delete_test.go 3.31 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
package server

import (
Michael Yang's avatar
Michael Yang committed
4
5
	"bytes"
	"encoding/json"
Michael Yang's avatar
Michael Yang committed
6
7
8
9
	"net/http"
	"path/filepath"
	"testing"

10
	"github.com/gin-gonic/gin"
Michael Yang's avatar
lint  
Michael Yang committed
11

Michael Yang's avatar
Michael Yang committed
12
	"github.com/ollama/ollama/api"
Michael Yang's avatar
Michael Yang committed
13
	"github.com/ollama/ollama/types/model"
Michael Yang's avatar
Michael Yang committed
14
15
16
)

func TestDelete(t *testing.T) {
17
18
	gin.SetMode(gin.TestMode)

Michael Yang's avatar
Michael Yang committed
19
20
	p := t.TempDir()
	t.Setenv("OLLAMA_MODELS", p)
21

Michael Yang's avatar
Michael Yang committed
22
23
	var s Server

24
	_, digest := createBinFile(t, nil, nil)
25
	w := createRequest(t, s.CreateHandler, api.CreateRequest{
26
27
		Name:  "test",
		Files: map[string]string{"test.gguf": digest},
Michael Yang's avatar
Michael Yang committed
28
29
30
31
32
33
	})

	if w.Code != http.StatusOK {
		t.Fatalf("expected status code 200, actual %d", w.Code)
	}

34
	w = createRequest(t, s.CreateHandler, api.CreateRequest{
35
36
37
		Name:     "test2",
		Files:    map[string]string{"test.gguf": digest},
		Template: "{{ .System }} {{ .Prompt }}",
Michael Yang's avatar
Michael Yang committed
38
39
40
41
42
43
44
45
46
47
48
49
	})

	if w.Code != http.StatusOK {
		t.Fatalf("expected status code 200, actual %d", w.Code)
	}

	checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
		filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
		filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
	})

	checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
50
51
52
		filepath.Join(p, "blobs", "sha256-136bf7c76bac2ec09d6617885507d37829e04b41acc47687d45e512b544e893a"),
		filepath.Join(p, "blobs", "sha256-6bcdb8859d417753645538d7bbfbd7ca91a3f0c191aef5379c53c05e86b669dd"),
		filepath.Join(p, "blobs", "sha256-89a2116c3a82d6a97f59f748d86ed4417214353fd178ee54df418fde32495fad"),
Michael Yang's avatar
Michael Yang committed
53
54
55
		filepath.Join(p, "blobs", "sha256-fe7ac77b725cda2ccad03f88a880ecdfd7a33192d6cae08fce2c0ee1455991ed"),
	})

56
	w = createRequest(t, s.DeleteHandler, api.DeleteRequest{Name: "test"})
Michael Yang's avatar
Michael Yang committed
57
58
59
60
61
62
63
64
65
66

	if w.Code != http.StatusOK {
		t.Fatalf("expected status code 200, actual %d", w.Code)
	}

	checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
		filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
	})

	checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
67
68
		filepath.Join(p, "blobs", "sha256-136bf7c76bac2ec09d6617885507d37829e04b41acc47687d45e512b544e893a"),
		filepath.Join(p, "blobs", "sha256-89a2116c3a82d6a97f59f748d86ed4417214353fd178ee54df418fde32495fad"),
Michael Yang's avatar
Michael Yang committed
69
70
71
		filepath.Join(p, "blobs", "sha256-fe7ac77b725cda2ccad03f88a880ecdfd7a33192d6cae08fce2c0ee1455991ed"),
	})

72
	w = createRequest(t, s.DeleteHandler, api.DeleteRequest{Name: "test2"})
Michael Yang's avatar
Michael Yang committed
73
74
75
76
77
78
79
80

	if w.Code != http.StatusOK {
		t.Fatalf("expected status code 200, actual %d", w.Code)
	}

	checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{})
	checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{})
}
Michael Yang's avatar
Michael Yang committed
81
82

func TestDeleteDuplicateLayers(t *testing.T) {
83
84
	gin.SetMode(gin.TestMode)

Michael Yang's avatar
Michael Yang committed
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
	p := t.TempDir()
	t.Setenv("OLLAMA_MODELS", p)
	var s Server

	n := model.ParseName("test")

	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(&ConfigV2{}); err != nil {
		t.Fatal(err)
	}

	config, err := NewLayer(&b, "application/vnd.docker.container.image.v1+json")
	if err != nil {
		t.Fatal(err)
	}

	// create a manifest with duplicate layers
102
	if err := WriteManifest(n, config, []Layer{config}); err != nil {
Michael Yang's avatar
Michael Yang committed
103
104
105
		t.Fatal(err)
	}

106
	w := createRequest(t, s.DeleteHandler, api.DeleteRequest{Name: "test"})
Michael Yang's avatar
Michael Yang committed
107
108
109
110
111
112
	if w.Code != http.StatusOK {
		t.Errorf("expected status code 200, actual %d", w.Code)
	}

	checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{})
}