1// Copyright 2023 Google LLC 2// 3// Use of this source code is governed by a BSD-style license that can be 4// found in the LICENSE file. 5 6package testutils 7 8import ( 9 "archive/zip" 10 "bytes" 11 "os" 12 "path/filepath" 13 "sort" 14 "testing" 15 16 "github.com/stretchr/testify/assert" 17 "github.com/stretchr/testify/require" 18 "go.skia.org/infra/task_driver/go/td" 19) 20 21// AssertStepNames flattens the names of the steps in the given report, then asserts that they 22// match the given list of strings. 23func AssertStepNames(t *testing.T, report *td.StepReport, expectedStepNames ...string) { 24 var actualStepNames []string 25 report.Recurse(func(sr *td.StepReport) bool { 26 if sr == report { 27 // The root step created by td.RunTestSteps() is always "fake-test-task". We can omit it. 28 return true 29 } 30 actualStepNames = append(actualStepNames, sr.Name) 31 return true 32 }) 33 assert.EqualValues(t, expectedStepNames, actualStepNames) 34} 35 36// MakeZIP creates a ZIP archive with the given contents, represented as a map from file paths to 37// file contents. 38func MakeZIP(t *testing.T, zipPath string, contents map[string]string) { 39 buf := &bytes.Buffer{} 40 zipWriter := zip.NewWriter(buf) 41 42 // Sort files before adding them to the archive for determinism. 43 var sortedFiles []string 44 for file := range contents { 45 sortedFiles = append(sortedFiles, file) 46 } 47 sort.Strings(sortedFiles) 48 49 for _, file := range sortedFiles { 50 fileWriter, err := zipWriter.Create(file) 51 require.NoError(t, err) 52 _, err = fileWriter.Write([]byte(contents[file])) 53 require.NoError(t, err) 54 } 55 56 require.NoError(t, zipWriter.Close()) 57 require.NoError(t, os.WriteFile(zipPath, buf.Bytes(), 0644)) 58} 59 60// PopulateDir populates a directory with the given contents, represented as a map from file paths 61// to file contents. 62func PopulateDir(t *testing.T, dirPath string, contents map[string]string) { 63 for filePath, fileContents := range contents { 64 require.NoError(t, os.WriteFile(filepath.Join(dirPath, filePath), []byte(fileContents), 0644)) 65 } 66} 67 68// MakeTempDirMockFn returns a function that can be used to mock os_steps.TempDir by setting the 69// os_steps.TempDirContextKey context key. It takes a list of directory paths, and returns the 70// first path on first call, then the second, and so on. If the function is called more times than 71// the number of directories, the test fails. 72func MakeTempDirMockFn(t *testing.T, dirs ...string) func(string, string) (string, error) { 73 require.NotEmpty(t, dirs) 74 i := 0 75 return func(string, string) (string, error) { 76 require.Less(t, i, len(dirs)) 77 dir := dirs[i] 78 i++ 79 return dir, nil 80 } 81} 82