github.com/alloyci/alloy-runner@v1.0.1-0.20180222164613-925503ccafd6/commands/helpers/artifacts_test.go (about)

     1  package helpers
     2  
     3  import (
     4  	"archive/zip"
     5  	"bytes"
     6  	"io"
     7  	"os"
     8  
     9  	"github.com/Sirupsen/logrus"
    10  
    11  	"gitlab.com/gitlab-org/gitlab-runner/common"
    12  )
    13  
    14  const artifactsTestArchivedFile = "archive_file"
    15  
    16  type testNetwork struct {
    17  	common.MockNetwork
    18  	downloadState  common.DownloadState
    19  	downloadCalled int
    20  	uploadState    common.UploadState
    21  	uploadCalled   int
    22  }
    23  
    24  func (m *testNetwork) DownloadArtifacts(config common.JobCredentials, artifactsFile string) common.DownloadState {
    25  	m.downloadCalled++
    26  
    27  	if m.downloadState == common.DownloadSucceeded {
    28  		file, err := os.Create(artifactsFile)
    29  		if err != nil {
    30  			logrus.Warningln(err)
    31  			return common.DownloadFailed
    32  		}
    33  		defer file.Close()
    34  
    35  		archive := zip.NewWriter(file)
    36  		archive.Create(artifactsTestArchivedFile)
    37  		archive.Close()
    38  	}
    39  	return m.downloadState
    40  }
    41  
    42  func (m *testNetwork) UploadRawArtifacts(config common.JobCredentials, reader io.Reader, baseName string, expireIn string) common.UploadState {
    43  	m.uploadCalled++
    44  
    45  	if m.uploadState == common.UploadSucceeded {
    46  		var buffer bytes.Buffer
    47  		io.Copy(&buffer, reader)
    48  		archive, err := zip.NewReader(bytes.NewReader(buffer.Bytes()), int64(buffer.Len()))
    49  		if err != nil {
    50  			logrus.Warningln(err)
    51  			return common.UploadForbidden
    52  		}
    53  
    54  		if len(archive.File) != 1 || archive.File[0].Name != artifactsTestArchivedFile {
    55  			logrus.Warningln("Invalid archive:", len(archive.File))
    56  			return common.UploadForbidden
    57  		}
    58  	}
    59  	return m.uploadState
    60  }