go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/artifacts/copy_test.go

// Copyright 2020 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 (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"

	"cloud.google.com/go/storage"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"google.golang.org/protobuf/types/known/structpb"

	"go.fuchsia.dev/infra/artifacts"
)

type mockDirectory struct {
	files       []string
	copiedFiles []string
	lock        sync.Mutex
}

func (d *mockDirectory) Object(path string) *storage.ObjectHandle {
	return nil
}

func (d *mockDirectory) CopyFile(ctx context.Context, src, dest string) (int64, error) {
	for _, file := range d.files {
		if file == src {
			d.lock.Lock()
			d.copiedFiles = append(d.copiedFiles, dest)
			d.lock.Unlock()
			return 1, nil
		}
	}
	return 0, os.ErrNotExist
}

func (d *mockDirectory) List(ctx context.Context, prefix string) ([]string, error) {
	var result []string
	for _, file := range d.files {
		if strings.HasPrefix(file, prefix) {
			result = append(result, file)
		}
	}
	if len(result) == 0 {
		return nil, &artifacts.ErrNothingMatchedPrefix{
			BucketName: "bucket",
			Prefix:     prefix,
		}
	}
	return result, nil
}

type mockArtifactsClient struct {
	files map[string][]string
	dir   *mockDirectory
}

func (a *mockArtifactsClient) GetBuildDir(bucket, buildID string) artifacts.Directory {
	a.dir = &mockDirectory{
		files: a.files[bucket],
	}
	return a.dir
}

func createSrcsFile(srcsFileContents []string) (*os.File, error) {
	f, err := os.CreateTemp("", "src-file")
	if err != nil {
		return nil, err
	}

	for _, src := range srcsFileContents {
		if _, err := fmt.Fprintf(f, "%s\n", src); err != nil {
			os.Remove(f.Name())
			return nil, err
		}
	}

	return f, nil
}

func TestExecute(t *testing.T) {
	archiveBucket := []string{
		"build-archive.tgz",
		"packages.tar.gz",
	}
	artifactsBucket := []string{
		"images/a",
		"images/b/c",
		"images/c",
		"packages/a",
		"packages/b/c",
	}
	storageContents := map[string][]string{
		"gcs_bucket":       archiveBucket,
		"artifacts_bucket": artifactsBucket,
	}
	artifactsCli := &mockArtifactsClient{
		files: storageContents,
	}
	buildID := "123"
	buildsCli := &mockBuildsClient{
		response: &buildbucketpb.Build{
			Id: 123,
			Output: &buildbucketpb.Build_Output{
				Properties: &structpb.Struct{
					Fields: map[string]*structpb.Value{
						"gcs_bucket": {
							Kind: &structpb.Value_StringValue{
								StringValue: "gcs_bucket",
							},
						},
						"artifact_gcs_bucket": {
							Kind: &structpb.Value_StringValue{
								StringValue: "artifacts_bucket",
							},
						},
					},
				},
			},
		},
	}
	tests := []struct {
		name             string
		source           string
		dest             string
		srcsFileContents []string
		expectedFiles    []string
		checkErr         func(t *testing.T, err error, stdout string)
	}{
		{
			name:          "copy directory from the artifacts bucket",
			source:        "packages",
			dest:          "output/packages",
			expectedFiles: []string{"output/packages/a", "output/packages/b/c"},
		},
		{
			name:          "copy file",
			source:        "images/a",
			dest:          "output/dest",
			expectedFiles: []string{"output/dest"},
		},
		{
			name:             "copy file from src-file",
			source:           "images/a",
			dest:             "output",
			srcsFileContents: []string{"images/a"},
			expectedFiles:    []string{"output"},
		},
		{
			name:             "copy multiple files from src-file",
			source:           "images",
			dest:             "output",
			srcsFileContents: []string{"images/a", "images/c"},
			expectedFiles:    []string{"output/a", "output/c"},
		},
		{
			name:             "copy from src-file with empty source",
			dest:             "output",
			srcsFileContents: []string{"images/a", "images/c"},
			expectedFiles:    []string{"output/images/a", "output/images/c"},
		},
		{
			name:   "copy src returns an error",
			source: "does-not-exist",
			dest:   "output/packages",
			checkErr: func(t *testing.T, err error, _ string) {
				var e *artifacts.ErrNothingMatchedPrefix
				if ok := errors.As(err, &e); !ok {
					t.Errorf("error should be ErrNothingMatchedPrefix, not %T", err)
				}

				if e.BucketName != "bucket" {
					t.Errorf("expected error bucket to be 'bucket', not %s", e.BucketName)
				}

				if e.Prefix != "does-not-exist" {
					t.Errorf("expected error bucket to be 'does-not-exist', not %s", e.Prefix)
				}
			},
		},
		{
			name:             "copy src-files returns an error",
			dest:             "output",
			srcsFileContents: []string{"does-not-exist", "images/a", "images/c"},
			checkErr: func(t *testing.T, err error, stdout string) {
				if !os.IsNotExist(err) {
					t.Errorf("expected error to be not found, not %s", err)
				}
				if !strings.Contains(stdout, "failed to download") {
					t.Errorf("expected log:\nfailed to download\nbut got:\n%+v", stdout)
				}
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			srcsFile := ""
			if tt.srcsFileContents != nil {
				f, err := createSrcsFile(tt.srcsFileContents)
				if err != nil {
					t.Errorf("failed to write src-file contents: %s", err)
				}
				defer os.Remove(f.Name())

				srcsFile = f.Name()
			}
			cmd := &CopyCommand{
				build:    buildID,
				source:   tt.source,
				dest:     tt.dest,
				srcsFile: srcsFile,
				j:        1,
				verbose:  true,
			}
			testStdout := &bytes.Buffer{}
			stdout = testStdout
			err := cmd.execute(context.Background(), buildsCli, artifactsCli)
			actualStdout := string(testStdout.Bytes())
			if tt.checkErr != nil {
				tt.checkErr(t, err, actualStdout)
				return
			} else if err != nil {
				t.Errorf("unexpected error: %s", err)
			}

			compare := func(expected []string, actual []string) {
				if len(expected) != len(actual) {
					t.Errorf("expected:\n%+v\nbut got:\n%+v", expected, actual)
				}
				fileMap := map[string]bool{}
				for _, file := range expected {
					fileMap[file] = true
				}
				for _, file := range actual {
					if _, ok := fileMap[file]; !ok {
						t.Errorf("expected:\n%+v\nbut got:\n%+v", expected, actual)
					}
				}
			}

			compare(tt.expectedFiles, artifactsCli.dir.copiedFiles)

			var expectedLogs []string
			for _, file := range tt.expectedFiles {
				relpath, err := filepath.Rel(tt.dest, file)
				if err != nil {
					t.Fatal(err)
				}
				source := filepath.Join(tt.source, relpath)
				expectedLogs = append(expectedLogs, fmt.Sprintf("%s (1 bytes) to %s", source, tt.dest))
			}
			numFiles := len(tt.expectedFiles)
			expectedLogs = append(expectedLogs, fmt.Sprintf("Num files: %d\nTotal bytes: %d", numFiles, numFiles))
			for _, log := range expectedLogs {
				if !strings.Contains(actualStdout, log) {
					t.Errorf("expected log:\n%+v\nbut got:\n%+v", log, actualStdout)
				}
			}
		})
	}
}