Update Unit Test.

Signed-off-by: yunkunrao <yunkunrao@gmail.com>

src/pkg/p2p/preheat/provider/client/http_client_test.go

@@ -44,6 +44,11 @@ func (suite *HTTPClientTestSuite) SetupSuite() {
 			return
 		}
 
+		// set http status code if needed
+		if r.URL.String() == "/statusCode" {
+			w.WriteHeader(http.StatusAlreadyReported)
+		}
+
 		w.Header().Add("Content-type", "application/json")
 		_, _ = w.Write([]byte("{}"))
 	}))
@@ -98,4 +103,7 @@ func (suite *HTTPClientTestSuite) TestPost() {
 	data, err := c.Post(suite.ts.URL, cred, []byte("{}"), map[string]string{"Accept": "application/json"})
 	suite.NoError(err, "post data")
 	suite.Equal("{}", string(data), "post json data")
+
+	data, err = c.Post(suite.ts.URL+"/statusCode", cred, []byte("{}"), map[string]string{"Accept": "application/json"})
+	suite.Error(err, "post data")
 }

src/pkg/p2p/preheat/provider/dragonfly_test.go

@@ -78,14 +78,33 @@ 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",
 		Tag:       "latest",
 		URL:       "https://harbor.com",
+		Digest:    "sha256:f3c97e3bd1e27393eb853a5c90b1132f2cda84336d5ba5d100c720dc98524c82",
 	})
 	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",
+	})
+	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")
 }
 
 // TestCheckProgress tests CheckProgress method.

src/pkg/p2p/preheat/provider/mock.go

@@ -25,6 +25,7 @@ import (
 )
 
 // This is a package to provide mock utilities.
+var preheatMap = make(map[string]struct{})
 
 // MockDragonflyProvider mocks a Dragonfly server.
 func MockDragonflyProvider() *httptest.Server {
@@ -57,6 +58,19 @@ func MockDragonflyProvider() *httptest.Server {
 				return
 			}
 
+			if image.ImageName == "" {
+				w.WriteHeader(http.StatusBadRequest)
+				return
+			}
+
+			if _, ok := preheatMap[image.Digest]; ok {
+				w.WriteHeader(http.StatusAlreadyReported)
+				_, _ = w.Write([]byte(`{"ID":""}`))
+				return
+			} else {
+				preheatMap[image.Digest] = struct{}{}
+			}
+
 			if image.Type == "image" &&
 				image.URL == "https://harbor.com" &&
 				image.ImageName == "busybox" &&