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
10
	"fmt"
	"net/http"
	"path/filepath"
	"testing"

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

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

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

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

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

25
	w := createRequest(t, s.CreateHandler, api.CreateRequest{
Michael Yang's avatar
Michael Yang committed
26
		Name:      "test",
27
		Modelfile: fmt.Sprintf("FROM %s", createBinFile(t, nil, nil)),
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{
Michael Yang's avatar
Michael Yang committed
35
		Name:      "test2",
36
		Modelfile: fmt.Sprintf("FROM %s\nTEMPLATE {{ .System }} {{ .Prompt }}", createBinFile(t, nil, nil)),
Michael Yang's avatar
Michael Yang committed
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
	})

	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{
		filepath.Join(p, "blobs", "sha256-8f2c2167d789c6b2302dff965160fa5029f6a24096d262c1cbb469f21a045382"),
		filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
		filepath.Join(p, "blobs", "sha256-ca239d7bd8ea90e4a5d2e6bf88f8d74a47b14336e73eb4e18bed4dd325018116"),
		filepath.Join(p, "blobs", "sha256-fe7ac77b725cda2ccad03f88a880ecdfd7a33192d6cae08fce2c0ee1455991ed"),
	})

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

	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{
		filepath.Join(p, "blobs", "sha256-8f2c2167d789c6b2302dff965160fa5029f6a24096d262c1cbb469f21a045382"),
		filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
		filepath.Join(p, "blobs", "sha256-fe7ac77b725cda2ccad03f88a880ecdfd7a33192d6cae08fce2c0ee1455991ed"),
	})

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

	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
80
81

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

Michael Yang's avatar
Michael Yang committed
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
	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
101
	if err := WriteManifest(n, config, []Layer{config}); err != nil {
Michael Yang's avatar
Michael Yang committed
102
103
104
		t.Fatal(err)
	}

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

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