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 (
	"context"
	"fmt"
	"strings"
	"sync"
	"testing"

	"cloud.google.com/go/storage"

	_struct "github.com/golang/protobuf/ptypes/struct"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"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) error {
	for _, file := range d.files {
		if file == src {
			d.lock.Lock()
			d.copiedFiles = append(d.copiedFiles, dest)
			d.lock.Unlock()
			return nil
		}
	}
	return fmt.Errorf("source doesn't exist")
}

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)
		}
	}
	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 TestExecute(t *testing.T) {
	archiveBucket := []string{
		"build-archive.tgz",
		"packages.tar.gz",
	}
	artifactsBucket := []string{
		"images/a",
		"images/b/c",
		"packages/a",
		"packages/b/c",
	}
	storageContents := map[string][]string{
		"gcs_bucket":       archiveBucket,
		"artifacts_bucket": artifactsBucket,
	}
	buildID := "123"
	buildsCli := &mockBuildsClient{
		response: &buildbucketpb.Build{
			Id: 123,
			Output: &buildbucketpb.Build_Output{
				Properties: &_struct.Struct{
					Fields: map[string]*_struct.Value{
						"gcs_bucket": {
							Kind: &_struct.Value_StringValue{
								StringValue: "gcs_bucket",
							},
						},
						"artifact_gcs_bucket": {
							Kind: &_struct.Value_StringValue{
								StringValue: "artifacts_bucket",
							},
						},
					},
				},
			},
		},
	}
	tests := []struct {
		name          string
		source        string
		dest          string
		expectedFiles []string
	}{
		{
			name:          "copy directory from the artifacts bucket",
			source:        "packages",
			dest:          "output/packages",
			expectedFiles: []string{"output/packages/a", "output/packages/b/c"},
		}, {
			name:          "copy archive",
			source:        "build-archive.tgz",
			dest:          "output/build-archive.tgz",
			expectedFiles: []string{"output/build-archive.tgz"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			artifactsCli := &mockArtifactsClient{
				files: storageContents,
			}
			cmd := &CopyCommand{
				build:  buildID,
				source: tt.source,
				dest:   tt.dest,
			}
			err := cmd.execute(context.Background(), buildsCli, artifactsCli)
			if err != nil {
				t.Errorf("unxpected error: %v", 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)
		})
	}
}