Skip to content

Commit

Permalink
test: add dragonfly p2p tests
Browse files Browse the repository at this point in the history 
Signed-off-by: Gaius <[email protected]>
  • Loading branch information
@gaius-qi
gaius-qi committed Sep 14, 2024
1 parent ba0b332 commit 64ed50b
Show file tree
Hide file tree
Showing 3 changed files with 123 additions and 150 deletions.
69 changes: 35 additions & 34 deletions src/pkg/p2p/preheat/provider/dragonfly.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,35 @@ func (dd *DragonflyDriver) Preheat(preheatingImage *PreheatImage) (*PreheatingSt
return nil, err
}

return &PreheatingStatus{
TaskID: resp.ID,
Status: provider.PreheatingStatusPending,
StartTime: resp.CreatedAt.Format(time.RFC3339),
FinishTime: resp.UpdatedAt.Format(time.RFC3339),
}, nil
}

// CheckProgress implements @Driver.CheckProgress.
func (dd *DragonflyDriver) CheckProgress(taskID string) (*PreheatingStatus, error) {
if dd.instance == nil {
return nil, errors.New("missing instance metadata")
}

if taskID == "" {
return nil, errors.New("no task ID")
}

url := fmt.Sprintf("%s%s/%s", strings.TrimSuffix(dd.instance.Endpoint, "/"), dragonflyJobPath, taskID)
data, err := client.GetHTTPClient(dd.instance.Insecure).Get(url, dd.getCred(), nil, nil)
if err != nil {
return nil, err
}

resp := &dragonflyJobResponse{}
if err := json.Unmarshal(data, resp); err != nil {
return nil, err
}

var (
successMessage string
errorMessage string
Expand All @@ -236,17 +265,17 @@ func (dd *DragonflyDriver) Preheat(preheatingImage *PreheatImage) (*PreheatingSt

var buffer bytes.Buffer
table := tablewriter.NewWriter(&buffer)
table.SetHeader([]string{"Hostname", "IP", "Cluster ID", "State"})
table.SetHeader([]string{"Hostname", "IP", "Cluster ID", "State", "Description"})
for _, jobState := range resp.Result.JobStates {
for _, result := range jobState.Results {
// Write the success tasks records to the table.
for _, successTask := range result.SuccessTasks {
table.Append([]string{successTask.Hostname, successTask.IP, string(result.SchedulerClusterID), dragonflyJobSuccessState})
table.Append([]string{successTask.Hostname, successTask.IP, fmt.Sprint(result.SchedulerClusterID), dragonflyJobSuccessState, ""})
}

// Write the failure tasks records to the table.
for _, failureTask := range result.FailureTasks {
table.Append([]string{failureTask.Hostname, failureTask.IP, string(result.SchedulerClusterID), dragonflyJobFailureState})
table.Append([]string{failureTask.Hostname, failureTask.IP, fmt.Sprint(result.SchedulerClusterID), dragonflyJobFailureState, failureTask.Description})
}
}
}
Expand All @@ -259,7 +288,9 @@ func (dd *DragonflyDriver) Preheat(preheatingImage *PreheatImage) (*PreheatingSt
errs = errors.Join(errs, errors.New(jobState.Error))
}

errorMessage = errs.Error()
if errs != nil {
errorMessage = errs.Error()
}
default:
state = provider.PreheatingStatusFail
errorMessage = fmt.Sprintf("unknown state: %s", resp.State)
Expand All @@ -275,36 +306,6 @@ func (dd *DragonflyDriver) Preheat(preheatingImage *PreheatImage) (*PreheatingSt
}, nil
}

// CheckProgress implements @Driver.CheckProgress.
func (dd *DragonflyDriver) CheckProgress(taskID string) (*PreheatingStatus, error) {
if dd.instance == nil {
return nil, errors.New("missing instance metadata")
}

if taskID == "" {
return nil, errors.New("no task ID")
}

url := fmt.Sprintf("%s%s/%s", strings.TrimSuffix(dd.instance.Endpoint, "/"), dragonflyJobPath, taskID)
data, err := client.GetHTTPClient(dd.instance.Insecure).Get(url, dd.getCred(), nil, nil)
if err != nil {
return nil, err
}

resp := &dragonflyJobResponse{}
if err := json.Unmarshal(data, resp); err != nil {
return nil, err
}

return &PreheatingStatus{
TaskID: resp.ID,
Status: resp.State,
Error: resp.Result.JobStates[0].Error,
StartTime: resp.CreatedAt.Format(time.RFC3339),
FinishTime: resp.UpdatedAt.Format(time.RFC3339),
}, nil
}

func (dd *DragonflyDriver) getCred() *auth.Credential {
return &auth.Credential{
Mode: dd.instance.AuthMode,
Expand Down
67 changes: 21 additions & 46 deletions src/pkg/p2p/preheat/provider/dragonfly_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,6 @@ func (suite *DragonflyTestSuite) TestGetHealth() {

// TestPreheat tests Preheat method.
func (suite *DragonflyTestSuite) TestPreheat() {
// preheat first time
st, err := suite.driver.Preheat(&PreheatImage{
Type: "image",
ImageName: "busybox",
Expand All @@ -89,56 +88,32 @@ func (suite *DragonflyTestSuite) TestPreheat() {
Scope: "single_peer",
})
require.NoError(suite.T(), err, "preheat image")
suite.Equal("dragonfly-id", st.TaskID, "preheat image result")

// preheat the same image second time
st, err = suite.driver.Preheat(&PreheatImage{
Type: "image",
ImageName: "busybox",
Tag: "latest",
URL: "https://harbor.com",
Digest: "sha256:f3c97e3bd1e27393eb853a5c90b1132f2cda84336d5ba5d100c720dc98524c82",
Scope: "all_peers",
})
require.NoError(suite.T(), err, "preheat image")
suite.Equal("", st.TaskID, "preheat image result")

// preheat image digest is empty
st, err = suite.driver.Preheat(&PreheatImage{
ImageName: "",
})
require.Error(suite.T(), err, "preheat image")
suite.Equal(provider.PreheatingStatusPending, st.Status, "preheat status")
suite.Equal("1", st.TaskID, "task id")
suite.NotEmptyf(st.StartTime, "start time")
suite.NotEmptyf(st.FinishTime, "finish time")
}

// TestCheckProgress tests CheckProgress method.
func (suite *DragonflyTestSuite) TestCheckProgress() {
st, err := suite.driver.CheckProgress("dragonfly-id")
require.NoError(suite.T(), err, "get preheat status")
st, err := suite.driver.CheckProgress("1")
require.NoError(suite.T(), err, "get image")
suite.Equal(provider.PreheatingStatusRunning, st.Status, "preheat status")
suite.Equal("1", st.TaskID, "task id")
suite.NotEmptyf(st.StartTime, "start time")
suite.NotEmptyf(st.FinishTime, "finish time")

st, err = suite.driver.CheckProgress("2")
require.NoError(suite.T(), err, "get image")
suite.Equal(provider.PreheatingStatusSuccess, st.Status, "preheat status")
suite.Equal("2", st.TaskID, "task id")
suite.NotEmptyf(st.StartTime, "start time")
suite.NotEmptyf(st.FinishTime, "finish time")

// preheat job exit but returns no id
st, err = suite.driver.CheckProgress("preheat-job-exist-with-no-id")
require.Error(suite.T(), err, "get preheat status")

// preheat job exit returns id but get info with that failed
st, err = suite.driver.CheckProgress("preheat-job-exist-with-id-1")
require.Error(suite.T(), err, "get preheat status")

// preheat job normal failed
st, err = suite.driver.CheckProgress("preheat-job-normal-failed")
require.NoError(suite.T(), err, "get preheat status")
st, err = suite.driver.CheckProgress("3")
require.NoError(suite.T(), err, "get image")
suite.Equal(provider.PreheatingStatusFail, st.Status, "preheat status")

// instance is empty
testDriver := &DragonflyDriver{}
st, err = testDriver.CheckProgress("")
require.Error(suite.T(), err, "get preheat status")

// preheat job with no task id
st, err = suite.driver.CheckProgress("")
require.Error(suite.T(), err, "get preheat status")

// preheat job with err json response
st, err = suite.driver.CheckProgress("preheat-job-err-body-json")
require.Error(suite.T(), err, "get preheat status")
suite.Equal("3", st.TaskID, "task id")
suite.NotEmptyf(st.StartTime, "start time")
suite.NotEmptyf(st.FinishTime, "finish time")
}
137 changes: 67 additions & 70 deletions src/pkg/p2p/preheat/provider/mock.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,10 @@ package provider

import (
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"time"

"github.com/goharbor/harbor/src/pkg/p2p/preheat/models/notification"
Expand All @@ -32,124 +32,121 @@ var preheatMap = make(map[string]struct{})
func MockDragonflyProvider() *httptest.Server {
return httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.RequestURI {
case healthCheckEndpoint:
case dragonflyHealthPath:
if r.Method != http.MethodGet {
w.WriteHeader(http.StatusNotImplemented)
return
}

w.WriteHeader(http.StatusOK)
case preheatEndpoint:
case dragonflyJobPath:
if r.Method != http.MethodPost {
w.WriteHeader(http.StatusNotImplemented)
return
}

data, err := io.ReadAll(r.Body)
var resp = &dragonflyJobResponse{
ID: "1",
State: dragonflyJobPendingState,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}

bytes, err := json.Marshal(resp)
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(err.Error()))
return
}

image := &PreheatImage{}
if err := json.Unmarshal(data, image); err != nil {
if _, err := w.Write(bytes); err != nil {
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(err.Error()))
return
}

if image.ImageName == "" {
w.WriteHeader(http.StatusBadRequest)
w.WriteHeader(http.StatusOK)
case fmt.Sprintf("%s/%s", dragonflyJobPath, "1"):
if r.Method != http.MethodGet {
w.WriteHeader(http.StatusNotImplemented)
return
}

if _, ok := preheatMap[image.Digest]; ok {
w.WriteHeader(http.StatusAlreadyReported)
_, _ = w.Write([]byte(`{"ID":""}`))
return
var resp = &dragonflyJobResponse{
ID: "1",
State: dragonflyJobPendingState,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}

preheatMap[image.Digest] = struct{}{}

if image.Type == "image" &&
image.URL == "https://harbor.com" &&
image.ImageName == "busybox" &&
image.Tag == "latest" {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{"ID":"dragonfly-id"}`))
bytes, err := json.Marshal(resp)
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(err.Error()))
return
}

w.WriteHeader(http.StatusBadRequest)
case strings.Replace(preheatTaskEndpoint, "{task_id}", "dragonfly-id", 1):
if r.Method != http.MethodGet {
w.WriteHeader(http.StatusNotImplemented)
if _, err := w.Write(bytes); err != nil {
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(err.Error()))
return
}
status := &dragonflyPreheatInfo{
ID: "dragonfly-id",
StartTime: time.Now().UTC().String(),
FinishTime: time.Now().Add(5 * time.Minute).UTC().String(),
Status: "SUCCESS",
}
bytes, _ := json.Marshal(status)
_, _ = w.Write(bytes)
case strings.Replace(preheatTaskEndpoint, "{task_id}", "preheat-job-exist-with-no-id", 1):

w.WriteHeader(http.StatusOK)
case fmt.Sprintf("%s/%s", dragonflyJobPath, "2"):
if r.Method != http.MethodGet {
w.WriteHeader(http.StatusNotImplemented)
return
}
status := &dragonflyPreheatInfo{
ID: "preheat-exist-with-no-id",
StartTime: time.Now().UTC().String(),
FinishTime: time.Now().Add(5 * time.Minute).UTC().String(),
Status: "FAILED",
ErrorMsg: "{\"Code\":208,\"Msg\":\"preheat task already exists, id:\"}",

var resp = &dragonflyJobResponse{
ID: "2",
State: dragonflyJobSuccessState,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
bytes, _ := json.Marshal(status)
_, _ = w.Write(bytes)
case strings.Replace(preheatTaskEndpoint, "{task_id}", "preheat-job-normal-failed", 1):
if r.Method != http.MethodGet {
w.WriteHeader(http.StatusNotImplemented)

bytes, err := json.Marshal(resp)
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(err.Error()))
return
}
status := &dragonflyPreheatInfo{
ID: "preheat-job-exist-with-id-1",
StartTime: time.Now().UTC().String(),
FinishTime: time.Now().Add(5 * time.Minute).UTC().String(),
Status: "FAILED",
ErrorMsg: "{\"Code\":208,\"Msg\":\"some msg\"}",

if _, err := w.Write(bytes); err != nil {
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(err.Error()))
return
}
bytes, _ := json.Marshal(status)
_, _ = w.Write(bytes)
case strings.Replace(preheatTaskEndpoint, "{task_id}", "preheat-job-exist-with-id-1", 1):

w.WriteHeader(http.StatusOK)
case fmt.Sprintf("%s/%s", dragonflyJobPath, "3"):
if r.Method != http.MethodGet {
w.WriteHeader(http.StatusNotImplemented)
return
}
status := &dragonflyPreheatInfo{
ID: "preheat-job-exist-with-id-1",
StartTime: time.Now().UTC().String(),
FinishTime: time.Now().Add(5 * time.Minute).UTC().String(),
Status: "FAILED",
ErrorMsg: "{\"Code\":208,\"Msg\":\"preheat task already exists, id:preheat-job-exist-with-id-1-1\"}",

var resp = &dragonflyJobResponse{
ID: "3",
State: dragonflyJobFailureState,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
bytes, _ := json.Marshal(status)
_, _ = w.Write(bytes)
case strings.Replace(preheatTaskEndpoint, "{task_id}", "preheat-job-exist-with-id-1-1", 1):
if r.Method != http.MethodGet {
w.WriteHeader(http.StatusNotImplemented)

bytes, err := json.Marshal(resp)
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(err.Error()))
return
}
w.WriteHeader(http.StatusInternalServerError)
case strings.Replace(preheatTaskEndpoint, "{task_id}", "preheat-job-err-body-json", 1):
if r.Method != http.MethodGet {
w.WriteHeader(http.StatusNotImplemented)

if _, err := w.Write(bytes); err != nil {
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(err.Error()))
return
}
bodyStr := "\"err body\""
_, _ = w.Write([]byte(bodyStr))

w.WriteHeader(http.StatusOK)
default:
w.WriteHeader(http.StatusNotImplemented)
}
Expand Down

0 comments on commit 64ed50b

Please sign in to comment.