pkg/coveragedb: test SaveMergeResult

1. Make interface testable.
2. Add Spanner interfaces.
3. Generate mocks for proxy interfaces.
4. Test SaveMergeResult.
5. Test MergeCSVWriteJSONL and coveragedb.SaveMergeResult integration.

Author: tarasmadan

Makefile

@@ -237,6 +237,7 @@ generate:
 generate_go: format_cpp
 	$(GO) generate ./executor ./pkg/ifuzz ./pkg/build ./pkg/rpcserver
 	$(GO) generate ./vm/proxyapp
+	$(GO) generate ./pkg/coveragedb
 
 generate_rpc:
 	flatc -o pkg/flatrpc --warnings-as-errors --gen-object-api --filename-suffix "" --go --gen-onefile --go-namespace flatrpc pkg/flatrpc/flatrpc.fbs

dashboard/app/api.go

@@ -26,6 +26,7 @@ import (
 	"github.com/google/syzkaller/pkg/asset"
 	"github.com/google/syzkaller/pkg/auth"
 	"github.com/google/syzkaller/pkg/coveragedb"
+	"github.com/google/syzkaller/pkg/coveragedb/spannerclient"
 	"github.com/google/syzkaller/pkg/debugtracer"
 	"github.com/google/syzkaller/pkg/email"
 	"github.com/google/syzkaller/pkg/gcs"
@@ -1944,15 +1945,21 @@ func apiCreateUploadURL(c context.Context, payload io.Reader) (interface{}, erro
 // Second+ records are coveragedb.MergedCoverageRecord.
 func apiSaveCoverage(c context.Context, payload io.Reader) (interface{}, error) {
 	descr := new(coveragedb.HistoryRecord)
-	if err := json.NewDecoder(payload).Decode(descr); err != nil {
-		return 0, fmt.Errorf("json.NewDecoder(dashapi.MergedCoverageDescription).Decode: %w", err)
+	jsonDec := json.NewDecoder(payload)
+	if err := jsonDec.Decode(descr); err != nil {
+		return 0, fmt.Errorf("json.NewDecoder(coveragedb.HistoryRecord).Decode: %w", err)
 	}
 	var sss []*subsystem.Subsystem
 	if service := getNsConfig(c, descr.Namespace).Subsystems.Service; service != nil {
 		sss = service.List()
 		log.Infof(c, "found %d subsystems for %s namespace", len(sss), descr.Namespace)
 	}
-	rowsCreated, err := coveragedb.SaveMergeResult(c, appengine.AppID(context.Background()), descr, payload, sss)
+	client, err := spannerclient.NewClient(c, appengine.AppID(context.Background()))
+	if err != nil {
+		return 0, fmt.Errorf("coveragedb.NewClient() failed: %s", err.Error())
+	}
+	defer client.Close()
+	rowsCreated, err := coveragedb.SaveMergeResult(c, client, descr, jsonDec, sss)
 	if err != nil {
 		log.Errorf(c, "error storing coverage for ns %s, date %s: %v",
 			descr.Namespace, descr.DateTo.String(), err)

dashboard/app/entities_spanner.go

@@ -11,6 +11,7 @@ import (
 	"cloud.google.com/go/civil"
 	"cloud.google.com/go/spanner"
 	"github.com/google/syzkaller/pkg/coveragedb"
+	"github.com/google/syzkaller/pkg/coveragedb/spannerclient"
 	"google.golang.org/api/iterator"
 )
 
@@ -25,7 +26,7 @@ type CoverageHistory struct {
 // MergedCoverage uses dates, not time.
 func MergedCoverage(ctx context.Context, ns, periodType string) (*CoverageHistory, error) {
 	projectID := os.Getenv("GOOGLE_CLOUD_PROJECT")
-	client, err := coveragedb.NewClient(ctx, projectID)
+	client, err := spannerclient.NewClient(ctx, projectID)
 	if err != nil {
 		return nil, fmt.Errorf("spanner.NewClient() failed: %s", err.Error())
 	}

pkg/cover/heatmap.go

@@ -14,6 +14,7 @@ import (
 
 	"cloud.google.com/go/spanner"
 	"github.com/google/syzkaller/pkg/coveragedb"
+	"github.com/google/syzkaller/pkg/coveragedb/spannerclient"
 	_ "github.com/google/syzkaller/pkg/subsystem/lists"
 	"golang.org/x/exp/maps"
 	"google.golang.org/api/iterator"
@@ -185,7 +186,7 @@ where
 
 func filesCoverageWithDetails(ctx context.Context, projectID, ns, subsystem string, timePeriods []coveragedb.TimePeriod,
 ) ([]*fileCoverageWithDetails, error) {
-	client, err := coveragedb.NewClient(ctx, projectID)
+	client, err := spannerclient.NewClient(ctx, projectID)
 	if err != nil {
 		return nil, fmt.Errorf("spanner.NewClient() failed: %s", err.Error())
 	}

pkg/coveragedb/spanner.go pkg/coveragedb/coveragedb.gopkg/coveragedb/spanner.go renamed to pkg/coveragedb/coveragedb.go

@@ -4,9 +4,9 @@
 package coveragedb
 
 import (
-	"bufio"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -15,6 +15,7 @@ import (
 
 	"cloud.google.com/go/civil"
 	"cloud.google.com/go/spanner"
+	"github.com/google/syzkaller/pkg/coveragedb/spannerclient"
 	"github.com/google/syzkaller/pkg/subsystem"
 	_ "github.com/google/syzkaller/pkg/subsystem/lists"
 	"github.com/google/uuid"
@@ -49,11 +50,6 @@ type HistoryRecord struct {
 	TotalRows int64
 }
 
-func NewClient(ctx context.Context, projectID string) (*spanner.Client, error) {
-	database := "projects/" + projectID + "/instances/syzbot/databases/coverage"
-	return spanner.NewClient(ctx, database)
-}
-
 type Coverage struct {
 	Instrumented      int64
 	Covered           int64
@@ -76,27 +72,26 @@ type MergedCoverageRecord struct {
 	FileData *Coverage
 }
 
-func SaveMergeResult(ctx context.Context, projectID string, descr *HistoryRecord, jsonl io.Reader,
+func SaveMergeResult(ctx context.Context, client spannerclient.SpannerClient, descr *HistoryRecord, dec *json.Decoder,
 	sss []*subsystem.Subsystem) (int, error) {
-	client, err := NewClient(ctx, projectID)
-	if err != nil {
-		return 0, fmt.Errorf("spanner.NewClient() failed: %s", err.Error())
-	}
-	defer client.Close()
 	var rowsCreated int
-
 	ssMatcher := subsystem.MakePathMatcher(sss)
 	ssCache := make(map[string][]string)
 
 	session := uuid.New().String()
 	mutations := []*spanner.Mutation{}
 
-	jsonlScanner := bufio.NewScanner(jsonl)
-
-	for jsonlScanner.Scan() {
-		var mcr MergedCoverageRecord
-		if err := json.Unmarshal([]byte(jsonlScanner.Text()), &mcr); err != nil {
-			return rowsCreated, fmt.Errorf("json.Unmarshal(MergedCoverageRecord): %w", err)
+	var mcr MergedCoverageRecord
+	for {
+		err := dec.Decode(&mcr)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return rowsCreated, fmt.Errorf("dec.Decode(MergedCoverageRecord): %w", err)
+		}
+		if mcr.FileData == nil {
+			return rowsCreated, errors.New("field MergedCoverageRecord.FileData can't be nil")
 		}
 		mutations = append(mutations, fileRecordMutation(session, &mcr))
 		subsystems := fileSubsystems(mcr.FilePath, ssMatcher, ssCache)
@@ -105,8 +100,8 @@ func SaveMergeResult(ctx context.Context, projectID string, descr *HistoryRecord
 		// This includes both explicit mutations of the fields (6 fields * 1k records = 6k mutations)
 		//   and implicit index mutations.
 		// We keep the number of records low enough for the number of explicit mutations * 10 does not exceed the limit.
-		if len(mutations) > 1000 {
-			if _, err = client.Apply(ctx, mutations); err != nil {
+		if len(mutations) >= 1000 {
+			if _, err := client.Apply(ctx, mutations); err != nil {
 				return rowsCreated, fmt.Errorf("failed to spanner.Apply(inserts): %s", err.Error())
 			}
 			rowsCreated += len(mutations)
@@ -115,9 +110,10 @@ func SaveMergeResult(ctx context.Context, projectID string, descr *HistoryRecord
 	}
 
 	mutations = append(mutations, historyMutation(session, descr))
-	if _, err = client.Apply(ctx, mutations); err != nil {
+	if _, err := client.Apply(ctx, mutations); err != nil {
 		return rowsCreated, fmt.Errorf("failed to spanner.Apply(inserts): %s", err.Error())
 	}
+	rowsCreated += len(mutations)
 	return rowsCreated, nil
 }
 
@@ -150,7 +146,7 @@ where
 func ReadLinesHitCount(ctx context.Context, ns, commit, file string, tp TimePeriod,
 ) (map[int]int, error) {
 	projectID := os.Getenv("GOOGLE_CLOUD_PROJECT")
-	client, err := NewClient(ctx, projectID)
+	client, err := spannerclient.NewClient(ctx, projectID)
 	if err != nil {
 		return nil, fmt.Errorf("spanner.NewClient: %w", err)
 	}
@@ -236,7 +232,7 @@ func fileSubsystems(filePath string, ssMatcher *subsystem.PathMatcher, ssCache m
 }
 
 func NsDataMerged(ctx context.Context, projectID, ns string) ([]TimePeriod, []int64, error) {
-	client, err := NewClient(ctx, projectID)
+	client, err := spannerclient.NewClient(ctx, projectID)
 	if err != nil {
 		return nil, nil, fmt.Errorf("spanner.NewClient() failed: %s", err.Error())
 	}
@@ -292,7 +288,7 @@ func NsDataMerged(ctx context.Context, projectID, ns string) ([]TimePeriod, []in
 // Returns the number of orphaned file entries successfully deleted.
 func DeleteGarbage(ctx context.Context) (int64, error) {
 	batchSize := 10_000
-	client, err := NewClient(ctx, os.Getenv("GOOGLE_CLOUD_PROJECT"))
+	client, err := spannerclient.NewClient(ctx, os.Getenv("GOOGLE_CLOUD_PROJECT"))
 	if err != nil {
 		return 0, fmt.Errorf("coveragedb.NewClient: %w", err)
 	}
@@ -339,7 +335,7 @@ func DeleteGarbage(ctx context.Context) (int64, error) {
 	return totalDeleted.Load(), nil
 }
 
-func goSpannerDelete(ctx context.Context, batch []spanner.Key, eg *errgroup.Group, client *spanner.Client,
+func goSpannerDelete(ctx context.Context, batch []spanner.Key, eg *errgroup.Group, client spannerclient.SpannerClient,
 	totalDeleted *atomic.Int64) {
 	ks := spanner.KeySetFromKeys(batch...)
 	ksSize := len(batch)

pkg/coveragedb/coveragedb_mock_test.go

@@ -0,0 +1,121 @@
+// Copyright 2024 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package coveragedb
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"strings"
+	"testing"
+	"time"
+
+	"cloud.google.com/go/spanner"
+	"github.com/google/syzkaller/pkg/coveragedb/mocks"
+	"github.com/google/syzkaller/pkg/subsystem"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+//go:generate ../../tools/mockery.sh --name SpannerClient -r
+
+type spannerMockTune func(*testing.T, *mocks.SpannerClient)
+
+func TestSaveMergeResult(t *testing.T) {
+	tests := []struct {
+		name     string
+		sss      []*subsystem.Subsystem
+		jsonl    io.Reader
+		descr    *HistoryRecord
+		mockTune spannerMockTune
+		wantErr  bool
+		wantRows int
+	}{
+		{
+			name:    "empty jsonl",
+			jsonl:   strings.NewReader(`{}`),
+			wantErr: true,
+		},
+		{
+			name:    "wrong jsonl content",
+			jsonl:   strings.NewReader(`{a}`),
+			wantErr: true,
+		},
+		{
+			name:     "1 record, Ok",
+			jsonl:    strings.NewReader(`{"FileData":{}}`),
+			descr:    &HistoryRecord{},
+			wantRows: 3, // 1 in files, 1 in file_subsystems and 1 in merge_history
+			mockTune: func(t *testing.T, m *mocks.SpannerClient) {
+				m.
+					On("Apply", mock.Anything, mock.Anything).
+					Return(time.Now(), nil).
+					Once()
+			},
+		},
+		{
+			name: "2 records, Ok",
+			jsonl: strings.NewReader(`	{"FileData":{}}
+																		{"FileData":{}}`),
+			descr:    &HistoryRecord{},
+			wantRows: 5,
+			mockTune: func(t *testing.T, m *mocks.SpannerClient) {
+				m.
+					On("Apply",
+						mock.Anything,
+						mock.MatchedBy(func(ms []*spanner.Mutation) bool {
+							// 2 in files, 2 in file_subsystems and 1 in merge_history
+							return len(ms) == 5
+						})).
+					Return(time.Now(), nil).
+					Once()
+			},
+		},
+		{
+			name:     "2k records, Ok",
+			jsonl:    strings.NewReader(strings.Repeat("{\"FileData\":{}}\n", 2000)),
+			descr:    &HistoryRecord{},
+			wantRows: 4001,
+			mockTune: func(t *testing.T, m *mocks.SpannerClient) {
+				m.
+					On("Apply",
+						mock.Anything,
+						mock.MatchedBy(func(ms []*spanner.Mutation) bool {
+							// 2k in files, 2k in file_subsystems
+							return len(ms) == 1000
+						})).
+					Return(time.Now(), nil).
+					Times(4).
+					On("Apply",
+						mock.Anything,
+						mock.MatchedBy(func(ms []*spanner.Mutation) bool {
+							// And 1 in merge_history.
+							return len(ms) == 1
+						})).
+					Return(time.Now(), nil).
+					Once()
+			},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+			spannerMock := mocks.NewSpannerClient(t)
+			if test.mockTune != nil {
+				test.mockTune(t, spannerMock)
+			}
+			gotRows, err := SaveMergeResult(
+				context.Background(),
+				spannerMock,
+				test.descr,
+				json.NewDecoder(test.jsonl), test.sss)
+			if test.wantErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+			assert.Equal(t, test.wantRows, gotRows)
+		})
+	}
+}