go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/update_test_durations/update_test.go

// Copyright 2022 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
)

func TestSplitTestsByBuilder(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		name             string
		options          durationFileOptions
		inputAndExpected func() (input []test, expected testDurationMap)
		expectedErr      error
	}{
		{
			name: "groups tests by builder",
			inputAndExpected: func() ([]test, testDurationMap) {
				fooTests := testsWithDurations("foo", 1, 1)
				barTests := testsWithDurations("bar", 1)

				allTests := appendAll(fooTests, barTests)

				expectedDurations := testDurationMap{
					"foo": fooTests,
					"bar": barTests,
				}
				return allTests, expectedDurations
			},
		},
		{
			name: "adds a default entry for each builder",
			options: durationFileOptions{
				includeDefaultTests: true,
			},
			inputAndExpected: func() ([]test, testDurationMap) {
				fooTests := testsWithDurations("foo", 1, 3)
				barTests := testsWithDurations("bar", 8, 10)

				allTests := appendAll(fooTests, barTests)

				expectedDurations := testDurationMap{
					"foo": append(fooTests, defaultEntryWithDuration(2)),
					"bar": append(barTests, defaultEntryWithDuration(9)),
				}
				return allTests, expectedDurations
			},
		},
		{
			name: "default entry duration is average of other entry durations",
			options: durationFileOptions{
				includeDefaultTests: true,
			},
			inputAndExpected: func() ([]test, testDurationMap) {
				// We want to use the average of medians rather than the median of
				// medians so that total expected shard durations are still accurate if
				// many new tests are added (assuming the new tests have a similar
				// distribution of durations to the existing tests).
				barTests := testsWithDurations("bar", 3, 30, 300)

				expectedDurations := testDurationMap{
					"bar": append(barTests, defaultEntryWithDuration(111)),
				}
				return barTests, expectedDurations
			},
		},
		{
			name: "weights default entry duration by run count",
			options: durationFileOptions{
				includeDefaultTests: true,
			},
			inputAndExpected: func() ([]test, testDurationMap) {
				allTests := []test{
					{
						Name:             "foo1",
						Builder:          "foo",
						MedianDurationMS: 150,
						Runs:             2,
					},
					{
						Name:             "foo1",
						Builder:          "foo",
						MedianDurationMS: 3,
						Runs:             1,
					},
				}

				expectedDurations := testDurationMap{
					// The default entry's duration should be the average of the other
					// entries' durations, weighted by run count.
					"foo": append(allTests, defaultEntryWithDuration(101)),
				}
				return allTests, expectedDurations
			},
		},
		{
			name: "adds a default builder",
			options: durationFileOptions{
				includeDefaultBuilder: true,
			},
			inputAndExpected: func() ([]test, testDurationMap) {
				sharedTest := test{Name: "shared", Runs: 2, MedianDurationMS: 5}
				fooSharedTest := test{Name: "shared", Runs: 1, MedianDurationMS: 5, Builder: "foo"}
				barSharedTest := test{Name: "shared", Runs: 1, MedianDurationMS: 5, Builder: "bar"}

				fooOnlyTests := testsWithDurations("foo", 1)
				fooTests := append(fooOnlyTests, fooSharedTest)
				barOnlyTests := testsWithDurations("bar", 3)
				barTests := append(barOnlyTests, barSharedTest)

				allTests := appendAll(fooTests, barTests)

				expectedDefaultTests := appendAll([]test{sharedTest}, fooOnlyTests, barOnlyTests)
				expectedDurations := testDurationMap{
					"foo":              fooTests,
					"bar":              barTests,
					defaultBuilderName: expectedDefaultTests,
				}
				return allTests, expectedDurations
			},
		},
		{
			name: "adds a default entry for the default builder",
			options: durationFileOptions{
				includeDefaultBuilder: true,
				includeDefaultTests:   true,
			},
			inputAndExpected: func() ([]test, testDurationMap) {
				fooTests := testsWithDurations("foo", 1)
				barTests := testsWithDurations("bar", 5)

				allTests := appendAll(fooTests, barTests)

				expectedDurations := testDurationMap{
					"foo":              append(fooTests, defaultEntryWithDuration(1)),
					"bar":              append(barTests, defaultEntryWithDuration(5)),
					defaultBuilderName: append(allTests, defaultEntryWithDuration(3)),
				}
				return allTests, expectedDurations
			},
		},
		{
			name: "returns no builders if input contains no tests",
			options: durationFileOptions{
				includeDefaultBuilder: true,
				includeDefaultTests:   true,
			},
			inputAndExpected: func() ([]test, testDurationMap) {
				var emptyTests []test
				var emptyDurations testDurationMap
				return emptyTests, emptyDurations
			},
		},
		{
			name: "returns error if all tests for a builder have zero runs",
			options: durationFileOptions{
				includeDefaultBuilder: true,
				includeDefaultTests:   true,
			},
			expectedErr: errZeroTotalRuns,
			inputAndExpected: func() ([]test, testDurationMap) {
				fooTests := testsWithDurations("foo", 1, 2, 3)
				for i := range fooTests {
					fooTests[i].Runs = 0
				}
				return fooTests, nil
			},
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			input, expected := tc.inputAndExpected()

			for builder, origTests := range expected {
				// Make a copy to avoid modifying the slices in the original, which may
				// share the same underlying array with the slice of input tests.
				tests := make([]test, len(origTests))
				copy(tests, origTests)
				expected[builder] = tests
				// It's very repetitive to specify the builder in each expected
				// test, so set it here based on the map key.
				for i := range tests {
					tests[i].Builder = builder
				}
			}

			actual, err := splitTestsByBuilder(input, tc.options)
			if !errors.Is(err, tc.expectedErr) {
				t.Fatalf("splitTestsByBuilder() returned wrong error, got: %v, wanted %v", err, tc.expectedErr)
			}

			opts := []cmp.Option{
				cmpopts.EquateEmpty(),
				// In production, duration files should always be sorted by name because
				// the Dremel query orders results by name. So there's no need for
				// splitTestsByBuilder to do sorting (except for the default duration
				// file), so we shouldn't care about ordering in its return value.
				cmpopts.SortSlices(func(a, b test) bool {
					return a.Name < b.Name
				}),
			}
			if diff := cmp.Diff(expected, actual, opts...); diff != "" {
				t.Errorf("splitTestsByBuilder() diff (-want +got):\n%s", diff)
			}
		})
	}
}

// testsWithDurations constructs a slice of tests with the given durations, in
// order, all have the same builder.
func testsWithDurations(builder string, durations ...int64) []test {
	var res []test
	for i, duration := range durations {
		res = append(res, test{
			Name:             fmt.Sprintf("%s-%d", builder, i),
			Builder:          builder,
			MedianDurationMS: duration,
			Runs:             1,
		})
	}
	return res
}

func defaultEntryWithDuration(medianDurationMS int64) test {
	return test{
		Name:             defaultTestName,
		MedianDurationMS: medianDurationMS,
	}
}

func appendAll(slices ...[]test) []test {
	var res []test
	for _, s := range slices {
		res = append(res, s...)
	}
	return res
}

func TestMarshalDurations(t *testing.T) {
	dir, err := os.MkdirTemp("/tmp", "test-durations-tests")
	if err != nil {
		t.Fatalf("Failed to create temporary directory: %v", err)
	}
	defer os.RemoveAll(dir)

	durations := testDurationMap{
		"foo": append(testsWithDurations("foo", 0, 1, 2), defaultEntryWithDuration(1)),
		"bar": append(testsWithDurations("bar", -3, 5), defaultEntryWithDuration(1)),
	}

	fileContents, err := marshalDurations(durations)
	if err != nil {
		t.Fatalf("marshalDurations() returned unexpected error: %v", err)
	}

	for builder, tests := range durations {
		// Use t.Run() so we can Fatalf() while still checking all duration files.
		t.Run(fmt.Sprintf("marshals durations for %s", builder), func(t *testing.T) {
			basename := fmt.Sprintf("%s.json", builder)
			b, ok := fileContents[basename]
			if !ok {
				t.Fatalf("File %s was not marshaled", basename)
			}
			var unmarshaledTests []test
			if err := json.Unmarshal(b, &unmarshaledTests); err != nil {
				t.Fatalf("Failed to unmarshal file %s: %v", basename, err)
			}

			var expectedTests []test
			for _, test := range tests {
				// The builder should not be included in the files, since it's in the
				// file name.
				test.Builder = ""
				expectedTests = append(expectedTests, test)
			}

			if diff := cmp.Diff(expectedTests, unmarshaledTests); diff != "" {
				t.Errorf("marshalDurations() diff in file %s (-want +got):\n%s", basename, diff)
			}
		})
	}
}

func TestOverlayFileContents(t *testing.T) {
	oldFiles := map[string][]byte{
		"foo": []byte("foo-old"),
		"bar": []byte("bar-old"),
	}

	newFiles := map[string][]byte{
		"bar": []byte("bar-new"),
		"baz": []byte("baz-new"),
	}

	expected := map[string][]byte{
		"foo": []byte("foo-old"),
		"bar": []byte("bar-new"),
		"baz": []byte("baz-new"),
	}

	result := overlayFileContents(oldFiles, newFiles)

	if diff := cmp.Diff(expected, result); diff != "" {
		t.Errorf("overlayFileContents() diff (-want +got):\n%s", diff)
	}
}