tools/botanist/fileutil_test.go

// Copyright 2019 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 botanist

import (
	"archive/tar"
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"sync"
	"testing"

	"go.fuchsia.dev/fuchsia/tools/net/tftp"
)

type retryTftpClient struct {
	failCount int
	failLimit int
	*tftp.ClientImpl
}

func (t *retryTftpClient) Read(ctx context.Context, path string) (*bytes.Reader, error) {
	if t.failCount < t.failLimit {
		t.failCount++
		return nil, fmt.Errorf("failed :(")
	}
	return bytes.NewReader([]byte{42, 42, 42}), nil
}

func TestFetchAndArchiveFile(t *testing.T) {
	ta, err := ioutil.TempFile("", t.Name())
	if err != nil {
		t.Fatalf("failed to create temp file: %s", err)
	}
	defer ta.Close()

	tw := tar.NewWriter(ta)
	defer tw.Close()

	client, err := tftp.NewClient(nil)
	if err != nil {
		t.Fatalf("failed to create tftp client: %s", err)
	}

	tftp := &retryTftpClient{
		failLimit:  1,
		ClientImpl: client,
	}

	if err := FetchAndArchiveFile(context.Background(), tftp, tw, "test/test", "test", &sync.Mutex{}); err != nil {
		t.Errorf("FetchAndArchive failed: %s", err)
	}
}