clair/database/pgsql/complex_test.go

package pgsql

import (
	"fmt"
	"math/rand"
	"runtime"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/coreos/clair/database"
	"github.com/coreos/clair/utils"
	"github.com/coreos/clair/utils/types"
	"github.com/pborman/uuid"
	"github.com/stretchr/testify/assert"
)

const (
	numVulnerabilities = 100
	numFeatureVersions = 100
)

func TestRaceAffects(t *testing.T) {
	datastore, err := OpenForTest("TestRaceAffects", false)
	if err != nil {
		t.Error(err)
		return
	}
	defer datastore.Close()

	// Insert the Feature on which we'll work.
	feature := database.Feature{
		Namespace: database.Namespace{Name: "TestRaceAffectsFeatureNamespace1"},
		Name:      "TestRaceAffecturesFeature1",
	}
	_, err = datastore.insertFeature(feature)
	if err != nil {
		t.Error(err)
		return
	}

	// Initialize random generator and enforce max procs.
	rand.Seed(time.Now().UnixNano())
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Generate FeatureVersions.
	featureVersions := make([]database.FeatureVersion, numFeatureVersions)
	for i := 0; i < numFeatureVersions; i++ {
		version := rand.Intn(numFeatureVersions)

		featureVersions[i] = database.FeatureVersion{
			Feature: feature,
			Version: types.NewVersionUnsafe(strconv.Itoa(version)),
		}
	}

	// Generate vulnerabilities.
	// They are mapped by fixed version, which will make verification really easy afterwards.
	vulnerabilities := make(map[int][]database.Vulnerability)
	for i := 0; i < numVulnerabilities; i++ {
		version := rand.Intn(numFeatureVersions) + 1

		// if _, ok := vulnerabilities[version]; !ok {
		//   vulnerabilities[version] = make([]database.Vulnerability)
		// }

		vulnerability := database.Vulnerability{
			Name:      uuid.New(),
			Namespace: feature.Namespace,
			FixedIn: []database.FeatureVersion{
				database.FeatureVersion{
					Feature: feature,
					Version: types.NewVersionUnsafe(strconv.Itoa(version)),
				},
			},
			Severity: types.Unknown,
		}

		vulnerabilities[version] = append(vulnerabilities[version], vulnerability)
	}

	// Insert featureversions and vulnerabilities in parallel.
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		for _, vulnerabilitiesM := range vulnerabilities {
			for _, vulnerability := range vulnerabilitiesM {
				err = datastore.InsertVulnerabilities([]database.Vulnerability{vulnerability})
				assert.Nil(t, err)
			}
		}
		fmt.Println("finished to insert vulnerabilities")
	}()

	go func() {
		defer wg.Done()
		for i := 0; i < len(featureVersions); i++ {
			featureVersions[i].ID, err = datastore.insertFeatureVersion(featureVersions[i])
			assert.Nil(t, err)
		}
		fmt.Println("finished to insert featureVersions")
	}()

	wg.Wait()

	// Verify consistency now.
	var actualAffectedNames []string
	var expectedAffectedNames []string

	for _, featureVersion := range featureVersions {
		featureVersionVersion, _ := strconv.Atoi(featureVersion.Version.String())

		// Get actual affects.
		rows, err := datastore.Query(getQuery("s_complextest_featureversion_affects"),
			featureVersion.ID)
		assert.Nil(t, err)
		defer rows.Close()

		var vulnName string
		for rows.Next() {
			err = rows.Scan(&vulnName)
			if !assert.Nil(t, err) {
				continue
			}
			actualAffectedNames = append(actualAffectedNames, vulnName)
		}
		if assert.Nil(t, rows.Err()) {
			rows.Close()
		}

		// Get expected affects.
		for i := numVulnerabilities; i > featureVersionVersion; i-- {
			for _, vulnerability := range vulnerabilities[i] {
				expectedAffectedNames = append(expectedAffectedNames, vulnerability.Name)
			}
		}

		assert.Len(t, utils.CompareStringLists(expectedAffectedNames, actualAffectedNames), 0)
		assert.Len(t, utils.CompareStringLists(actualAffectedNames, expectedAffectedNames), 0)
	}

	// TODO(Quentin-M): May be worth having a test for updates as well.
}