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 committed Sep 14, 2024
1 parent ba0b332 commit f6df0b3
Show file tree
Hide file tree
Showing 3 changed files with 123 additions and 156 deletions.
75 changes: 35 additions & 40 deletions src/pkg/p2p/preheat/provider/dragonfly.go
Original file line number Diff line number Diff line change
Expand Up @@ -116,9 +116,6 @@ type dragonflyJobResponse struct {
// SuccessTasks is the success tasks.
SuccessTasks []*struct {

// TaskID is the task id.
TaskID string `json:"task_id"`

// Hostname is the hostname of the task.
Hostname string `json:"hostname"`

Expand All @@ -129,9 +126,6 @@ type dragonflyJobResponse struct {
// FailureTasks is the failure tasks.
FailureTasks []*struct {

// TaskID is the task id.
TaskID string `json:"task_id"`

// Hostname is the hostname of the task.
Hostname string `json:"hostname"`

Expand Down Expand Up @@ -223,6 +217,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 +259,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 +282,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 +300,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
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")
}
Loading

0 comments on commit f6df0b3

Please sign in to comment.