1 files changed, 98 insertions, 0 deletions
diff --git a/pkg/corpus/corpus_test.go b/pkg/corpus/corpus_test.go
new file mode 100644
index 000000000..cce537087
--- /dev/null
+++ b/pkg/corpus/corpus_test.go
@@ -0,0 +1,98 @@
+// 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 corpus
+
+import (
+	"context"
+	"math/rand"
+	"testing"
+
+	"github.com/google/syzkaller/pkg/signal"
+	"github.com/google/syzkaller/prog"
+	"github.com/google/syzkaller/sys/targets"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCorpusOperation(t *testing.T) {
+	// Basic corpus functionality.
+	target := getTarget(t, targets.TestOS, targets.TestArch64)
+	ch := make(chan NewItemEvent)
+	corpus := NewMonitoredCorpus(context.Background(), ch)
+
+	// First program is saved.
+	rs := rand.NewSource(0)
+	inp1 := generateInput(target, rs, 5, 5)
+	go corpus.Save(inp1)
+	event := <-ch
+	progData := inp1.Prog.Serialize()
+	assert.Equal(t, progData, event.ProgData)
+	assert.Equal(t, false, event.Exists)
+
+	// Second program is saved for every its call.
+	inp2 := generateInput(target, rs, 5, 5)
+	progData = inp2.Prog.Serialize()
+	for i := 0; i < 5; i++ {
+		inp2.Call = i
+		go corpus.Save(inp2)
+		event := <-ch
+		assert.Equal(t, progData, event.ProgData)
+		assert.Equal(t, i != 0, event.Exists)
+	}
+
+	// Verify that we can query corpus items.
+	items := corpus.Items()
+	assert.Len(t, items, 2)
+	for _, item := range items {
+		assert.Equal(t, item, corpus.Item(item.Sig))
+	}
+
+	// Verify the total signal.
+	assert.Len(t, corpus.Signal(), 5)
+
+	corpus.Minimize(true)
+}
+
+func TestCorpusSaveConcurrency(t *testing.T) {
+	target := getTarget(t, targets.TestOS, targets.TestArch64)
+	corpus := NewCorpus(context.Background())
+
+	const (
+		routines = 10
+		iters    = 100
+	)
+
+	for i := 0; i < routines; i++ {
+		go func() {
+			rs := rand.NewSource(0)
+			r := rand.New(rs)
+			for it := 0; it < iters; it++ {
+				inp := generateInput(target, rs, 10, it)
+				corpus.Save(inp)
+				corpus.ChooseProgram(r).Clone()
+			}
+		}()
+	}
+}
+
+func generateInput(target *prog.Target, rs rand.Source, ncalls, sizeSig int) NewInput {
+	p := target.Generate(rs, ncalls, target.DefaultChoiceTable())
+	var raw []uint32
+	for i := 1; i <= sizeSig; i++ {
+		raw = append(raw, uint32(i))
+	}
+	return NewInput{
+		Prog:   p,
+		Call:   int(rs.Int63() % int64(len(p.Calls))),
+		Signal: signal.FromRaw(raw, 0),
+	}
+}
+
+func getTarget(t *testing.T, os, arch string) *prog.Target {
+	t.Parallel()
+	target, err := prog.GetTarget(os, arch)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return target
+}