pkg/syz-ci: test func UploadCoverJSONLToGCS

Author: tarasmadan

syz-ci/manager.go

@@ -866,7 +866,7 @@ func (mgr *Manager) uploadCoverReport() error {
 	}
 	defer resp.Body.Close()
 	if directUpload {
-		return uploadFile(mgr.cfg.CoverUploadPath, mgr.name+".html", resp.Body, mgr.cfg.PublishGCS)
+		return uploadFile(context.Background(), mgr.cfg.CoverUploadPath, mgr.name+".html", resp.Body, mgr.cfg.PublishGCS)
 	}
 	// Upload via the asset storage.
 	newAsset, err := mgr.storage.UploadBuildAsset(resp.Body, mgr.name+".html",
@@ -881,8 +881,8 @@ func (mgr *Manager) uploadCoverReport() error {
 	return nil
 }
 
-func (mgr *Manager) uploadCoverJSONLToGCS(mgrSrc, gcsDest string, curTime time.Time, publish, compress bool,
-	f func(io.Writer, *json.Decoder) error) error {
+func (mgr *Manager) uploadCoverJSONLToGCS(ctx context.Context, mgrSrc, gcsDest string, curTime time.Time,
+	publish, compress bool, f func(io.Writer, *json.Decoder) error) error {
 	if !mgr.managercfg.Cover || gcsDest == "" {
 		return nil
 	}
@@ -932,7 +932,7 @@ func (mgr *Manager) uploadCoverJSONLToGCS(mgrSrc, gcsDest string, curTime time.T
 		mgr.mgrcfg.DashboardClient,
 		mgr.name, curTime.Format(time.DateOnly),
 		curTime.Hour(), curTime.Minute())
-	if err := uploadFile(gcsDest, fileName, pr, publish); err != nil {
+	if err := uploadFile(ctx, gcsDest, fileName, pr, publish); err != nil {
 		return fmt.Errorf("failed to uploadFileGCS(): %w", err)
 	}
 	return nil
@@ -943,7 +943,8 @@ func (mgr *Manager) uploadCoverStat(fuzzingMinutes int) error {
 	// In the syz-ci context report generation won't be used after this point,
 	// so tell manager to flush report generator.
 	curTime := time.Now()
-	if err := mgr.uploadCoverJSONLToGCS("/cover?jsonl=1&flush=1",
+	if err := mgr.uploadCoverJSONLToGCS(context.Background(),
+		"/cover?jsonl=1&flush=1",
 		mgr.cfg.CoverPipelinePath,
 		curTime,
 		false,
@@ -974,7 +975,8 @@ func (mgr *Manager) uploadCoverStat(fuzzingMinutes int) error {
 }
 
 func (mgr *Manager) uploadProgramsWithCoverage() error {
-	if err := mgr.uploadCoverJSONLToGCS("/coverprogs?jsonl=1",
+	if err := mgr.uploadCoverJSONLToGCS(context.Background(),
+		"/coverprogs?jsonl=1",
 		mgr.cfg.CoverProgramsPath,
 		time.Now(),
 		mgr.cfg.PublishGCS,
@@ -1005,7 +1007,7 @@ func (mgr *Manager) uploadCorpus() error {
 		return err
 	}
 	defer f.Close()
-	return uploadFile(mgr.cfg.CorpusUploadPath, mgr.name+"-corpus.db", f, mgr.cfg.PublishGCS)
+	return uploadFile(context.Background(), mgr.cfg.CorpusUploadPath, mgr.name+"-corpus.db", f, mgr.cfg.PublishGCS)
 }
 
 func (mgr *Manager) uploadBenchData() error {
@@ -1023,15 +1025,15 @@ func (mgr *Manager) uploadBenchData() error {
 		return fmt.Errorf("failed to open bench file: %w", err)
 	}
 	defer f.Close()
-	err = uploadFile(mgr.cfg.BenchUploadPath+"/"+mgr.name,
+	err = uploadFile(context.Background(), mgr.cfg.BenchUploadPath+"/"+mgr.name,
 		mgr.lastRestarted.Format("2006-01-02_15h.json"), f, false)
 	if err != nil {
 		return fmt.Errorf("failed to upload the bench file: %w", err)
 	}
 	return nil
 }
 
-func uploadFile(dstPath, name string, file io.Reader, publish bool) error {
+func uploadFile(ctx context.Context, dstPath, name string, file io.Reader, publish bool) error {
 	URL, err := url.Parse(dstPath)
 	if err != nil {
 		return fmt.Errorf("failed to parse upload path: %w", err)
@@ -1043,7 +1045,7 @@ func uploadFile(dstPath, name string, file io.Reader, publish bool) error {
 		strings.HasPrefix(URLStr, "https://") {
 		return uploadFileHTTPPut(URLStr, file)
 	}
-	return gcs.UploadFile(context.Background(), file, URLStr, publish)
+	return gcs.UploadFile(ctx, file, URLStr, publish)
 }
 
 func uploadFileHTTPPut(URL string, file io.Reader) error {

syz-ci/manager_test.go

@@ -4,10 +4,23 @@
 package main
 
 import (
+	"bytes"
+	"compress/gzip"
+	"context"
+	"encoding/json"
 	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/google/syzkaller/dashboard/dashapi"
+	"github.com/google/syzkaller/pkg/cover"
+	"github.com/google/syzkaller/pkg/gcs"
+	gcsmocks "github.com/google/syzkaller/pkg/gcs/mocks"
+	"github.com/google/syzkaller/pkg/mgrconfig"
 	"github.com/google/syzkaller/pkg/vcs"
 	"github.com/google/syzkaller/sys/targets"
 	"github.com/stretchr/testify/assert"
@@ -106,3 +119,130 @@ Reported-by: [email protected]`,
 	assert.Equal(t, commit.Title, "title with fix")
 	assert.ElementsMatch(t, commit.BugIDs, []string{"abcd000"})
 }
+
+func TestUploadCoverJSONLToGCS(t *testing.T) {
+	tests := []struct {
+		name string
+
+		inputJSONL string
+		inputTime  time.Time
+
+		inputCompress bool
+		inputPublish  bool
+
+		wantGCSFileName    string
+		wantGCSFileContent string
+		wantCompressed     bool
+		wantPublish        bool
+		wantError          string
+	}{
+		{
+			name:               "upload single object",
+			inputJSONL:         "{}",
+			inputTime:          time.Time{},
+			wantGCSFileName:    "test-bucket/test-namespace/mgr-name-0001-01-01-0-0.jsonl",
+			wantGCSFileContent: "{}\n",
+		},
+		{
+			name:               "upload single object, compress",
+			inputJSONL:         "{}",
+			inputTime:          time.Time{},
+			inputCompress:      true,
+			wantGCSFileName:    "test-bucket/test-namespace/mgr-name-0001-01-01-0-0.jsonl",
+			wantGCSFileContent: "{}\n",
+			wantCompressed:     true,
+		},
+		{
+			name:               "upload single object, publish",
+			inputJSONL:         "{}",
+			inputTime:          time.Time{},
+			inputPublish:       true,
+			wantGCSFileName:    "test-bucket/test-namespace/mgr-name-0001-01-01-0-0.jsonl",
+			wantGCSFileContent: "{}\n",
+			wantPublish:        true,
+		},
+		{
+			name:            "upload single object, error",
+			inputJSONL:      "{",
+			inputTime:       time.Time{},
+			wantGCSFileName: "test-bucket/test-namespace/mgr-name-0001-01-01-0-0.jsonl",
+			wantError:       "failed to uploadFileGCS(): io.Copy: callback: cover.ProgramCoverage: unexpected EOF",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+				w.Write([]byte(test.inputJSONL))
+			}))
+			defer httpServer.Close()
+
+			testSetverAddrPort, _ := strings.CutPrefix(httpServer.URL, "http://")
+			mgr := Manager{
+				name: "mgr-name",
+				managercfg: &mgrconfig.Config{
+					HTTP:  testSetverAddrPort,
+					Cover: true,
+				},
+				mgrcfg: &ManagerConfig{
+					DashboardClient: "test-namespace",
+				},
+			}
+
+			gcsMock := gcsmocks.NewClient(t)
+			gotBytes := mockWriteCloser{}
+
+			gcsMock.On("FileWriter", test.wantGCSFileName).
+				Return(&gotBytes, nil).Once()
+			gcsMock.On("Close").Return().Once()
+			if test.wantPublish {
+				gcsMock.On("Publish", test.wantGCSFileName).
+					Return(nil).Once()
+			}
+			ctx := gcs.SetGCSClient(context.Background(), gcsMock)
+			err := mgr.uploadCoverJSONLToGCS(ctx,
+				"/teststream&jsonl=1",
+				"gs://test-bucket",
+				time.Time{}, test.inputPublish, test.inputCompress, func(w io.Writer, dec *json.Decoder) error {
+					var v any
+					if err := dec.Decode(&v); err != nil {
+						return fmt.Errorf("cover.ProgramCoverage: %w", err)
+					}
+					if err := cover.WriteJSLine(w, &v); err != nil {
+						return fmt.Errorf("cover.WriteJSLine: %w", err)
+					}
+					return nil
+				})
+			if test.wantError != "" {
+				assert.Equal(t, test.wantError, err.Error())
+			} else {
+				assert.NoError(t, err)
+			}
+			assert.Equal(t, 1, gotBytes.closedTimes)
+			if test.wantCompressed {
+				gzReader, err := gzip.NewReader(&gotBytes.buf)
+				assert.NoError(t, err)
+				defer gzReader.Close()
+				plainBytes := mockWriteCloser{}
+				_, err = io.Copy(&plainBytes, gzReader)
+				assert.NoError(t, err)
+				gotBytes = plainBytes
+			}
+			assert.Equal(t, test.wantGCSFileContent, gotBytes.buf.String())
+		})
+	}
+}
+
+type mockWriteCloser struct {
+	buf         bytes.Buffer
+	closedTimes int
+}
+
+func (m *mockWriteCloser) Write(p []byte) (n int, err error) {
+	return m.buf.Write(p)
+}
+
+func (m *mockWriteCloser) Close() error {
+	m.closedTimes++
+	return nil
+}