From 849427636eb5671636f5e1b50cd6988426b175eb Mon Sep 17 00:00:00 2001
From: Gaius <gaius.qi@gmail.com>
Date: Thu, 12 Sep 2024 21:17:45 +0800
Subject: [PATCH] test: add dragonfly p2p tests

Signed-off-by: Gaius <gaius.qi@gmail.com>
---
 src/pkg/p2p/preheat/instance/manager.go   | 14 +++++++++++++-
 src/pkg/p2p/preheat/job.go                |  5 +++++
 src/pkg/p2p/preheat/provider/dragonfly.go | 20 ++++++++------------
 src/pkg/p2p/preheat/provider/mock.go      |  8 ++++----
 4 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/src/pkg/p2p/preheat/instance/manager.go b/src/pkg/p2p/preheat/instance/manager.go
index acb43b0c9e1..3f4418a1013 100644
--- a/src/pkg/p2p/preheat/instance/manager.go
+++ b/src/pkg/p2p/preheat/instance/manager.go
@@ -16,6 +16,7 @@ package instance
 
 import (
 	"context"
+	"encoding/json"
 
 	"github.com/goharbor/harbor/src/lib/q"
 	dao "github.com/goharbor/harbor/src/pkg/p2p/preheat/dao/instance"
@@ -114,7 +115,18 @@ func (dm *manager) Update(ctx context.Context, inst *provider.Instance, props ..
 
 // Get implements @Manager.Get
 func (dm *manager) Get(ctx context.Context, id int64) (*provider.Instance, error) {
-	return dm.dao.Get(ctx, id)
+	ins, err := dm.dao.Get(ctx, id)
+	if err != nil {
+		return nil, err
+	}
+	// mapping auth data to auth info.
+	if len(ins.AuthData) > 0 {
+		if err := json.Unmarshal([]byte(ins.AuthData), &ins.AuthInfo); err != nil {
+			return nil, err
+		}
+	}
+
+	return ins, nil
 }
 
 // Get implements @Manager.GetByName
diff --git a/src/pkg/p2p/preheat/job.go b/src/pkg/p2p/preheat/job.go
index 7321260b196..ee68a6272d3 100644
--- a/src/pkg/p2p/preheat/job.go
+++ b/src/pkg/p2p/preheat/job.go
@@ -191,6 +191,11 @@ func (j *Job) Run(ctx job.Context, params job.Parameters) error {
 				return preheatJobRunningError(errors.Errorf("preheat failed: %s", s))
 			case provider.PreheatingStatusSuccess:
 				// Finished
+				// log the message if received message from provider.
+				if s.Message != "" {
+					myLogger.Infof("Preheat job finished, message from provider: \n%s", s.Message)
+				}
+
 				return nil
 			default:
 				// do nothing, check again
diff --git a/src/pkg/p2p/preheat/provider/dragonfly.go b/src/pkg/p2p/preheat/provider/dragonfly.go
index edb2a162fee..42f50235f64 100644
--- a/src/pkg/p2p/preheat/provider/dragonfly.go
+++ b/src/pkg/p2p/preheat/provider/dragonfly.go
@@ -17,12 +17,12 @@ package provider
 import (
 	"bytes"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"strings"
 	"time"
 
 	"github.com/goharbor/harbor/src/lib"
+	"github.com/goharbor/harbor/src/lib/errors"
 	"github.com/goharbor/harbor/src/pkg/p2p/preheat/models/provider"
 	"github.com/goharbor/harbor/src/pkg/p2p/preheat/provider/auth"
 	"github.com/goharbor/harbor/src/pkg/p2p/preheat/provider/client"
@@ -90,7 +90,7 @@ type dragonflyCreateJobRequestArgs struct {
 
 type dragonflyJobResponse struct {
 	// ID is the job id.
-	ID string `json:"id"`
+	ID int `json:"id"`
 
 	// CreatedAt is the job created time.
 	CreatedAt time.Time `json:"created_at"`
@@ -224,7 +224,7 @@ func (dd *DragonflyDriver) Preheat(preheatingImage *PreheatImage) (*PreheatingSt
 	}
 
 	return &PreheatingStatus{
-		TaskID:     resp.ID,
+		TaskID:     fmt.Sprintf("%d", resp.ID),
 		Status:     provider.PreheatingStatusPending,
 		StartTime:  resp.CreatedAt.Format(time.RFC3339),
 		FinishTime: resp.UpdatedAt.Format(time.RFC3339),
@@ -283,18 +283,14 @@ func (dd *DragonflyDriver) CheckProgress(taskID string) (*PreheatingStatus, erro
 		table.Render()
 		successMessage = buffer.String()
 	case dragonflyJobFailureState:
+		var errs errors.Errors
 		state = provider.PreheatingStatusFail
-
-		var buffer bytes.Buffer
-		table := tablewriter.NewWriter(&buffer)
-		table.SetHeader([]string{"Error Message"})
 		for _, jobState := range resp.Result.JobStates {
-			table.Append([]string{jobState.Error})
+			errs = append(errs, errors.New(jobState.Error))
 		}
 
-		table.Render()
-		if len(resp.Result.JobStates) > 0 {
-			errorMessage = buffer.String()
+		if len(errs) > 0 {
+			errorMessage = errs.Error()
 		}
 	default:
 		state = provider.PreheatingStatusFail
@@ -302,7 +298,7 @@ func (dd *DragonflyDriver) CheckProgress(taskID string) (*PreheatingStatus, erro
 	}
 
 	return &PreheatingStatus{
-		TaskID:     resp.ID,
+		TaskID:     fmt.Sprintf("%d", resp.ID),
 		Status:     state,
 		Message:    successMessage,
 		Error:      errorMessage,
diff --git a/src/pkg/p2p/preheat/provider/mock.go b/src/pkg/p2p/preheat/provider/mock.go
index aa38f894c43..843f7872b00 100644
--- a/src/pkg/p2p/preheat/provider/mock.go
+++ b/src/pkg/p2p/preheat/provider/mock.go
@@ -46,7 +46,7 @@ func MockDragonflyProvider() *httptest.Server {
 			}
 
 			var resp = &dragonflyJobResponse{
-				ID:        "1",
+				ID:        1,
 				State:     dragonflyJobPendingState,
 				CreatedAt: time.Now(),
 				UpdatedAt: time.Now(),
@@ -73,7 +73,7 @@ func MockDragonflyProvider() *httptest.Server {
 			}
 
 			var resp = &dragonflyJobResponse{
-				ID:        "1",
+				ID:        1,
 				State:     dragonflyJobPendingState,
 				CreatedAt: time.Now(),
 				UpdatedAt: time.Now(),
@@ -100,7 +100,7 @@ func MockDragonflyProvider() *httptest.Server {
 			}
 
 			var resp = &dragonflyJobResponse{
-				ID:        "2",
+				ID:        2,
 				State:     dragonflyJobSuccessState,
 				CreatedAt: time.Now(),
 				UpdatedAt: time.Now(),
@@ -127,7 +127,7 @@ func MockDragonflyProvider() *httptest.Server {
 			}
 
 			var resp = &dragonflyJobResponse{
-				ID:        "3",
+				ID:        3,
 				State:     dragonflyJobFailureState,
 				CreatedAt: time.Now(),
 				UpdatedAt: time.Now(),