| package gcs |
| |
| import ( |
| "archive/tar" |
| "bytes" |
| "compress/gzip" |
| "context" |
| "fmt" |
| "io" |
| "io/ioutil" |
| "os" |
| "path/filepath" |
| "strings" |
| |
| "cloud.google.com/go/storage" |
| assert "github.com/stretchr/testify/require" |
| "go.skia.org/infra/go/httputils" |
| "go.skia.org/infra/go/util" |
| "google.golang.org/api/option" |
| ) |
| |
| const ( |
| // GCS bucket where we store test data. Add a folder to this bucket |
| // with the tests for a particular component. |
| TEST_DATA_BUCKET = "skia-infra-testdata" |
| ) |
| |
| func getStorangeItem(bucket, gsPath string) (*storage.Reader, error) { |
| storageClient, err := storage.NewClient(context.Background(), option.WithHTTPClient(httputils.NewTimeoutClient())) |
| if err != nil { |
| return nil, err |
| } |
| |
| return storageClient.Bucket(bucket).Object(gsPath).NewReader(context.Background()) |
| } |
| |
| // DownloadTestDataFile downloads a file with test data from Google Storage. |
| // The uriPath identifies what to download from the test bucket in GCS. |
| // The content must be publicly accessible. |
| // The file will be downloaded and stored at provided target |
| // path (regardless of what the original name is). |
| // If the the uri ends with '.gz' it will be transparently unzipped. |
| func DownloadTestDataFile(t assert.TestingT, bucket, gsPath, targetPath string) error { |
| dir, _ := filepath.Split(targetPath) |
| if err := os.MkdirAll(dir, 0755); err != nil { |
| return err |
| } |
| |
| arch, err := getStorangeItem(bucket, gsPath) |
| if err != nil { |
| return fmt.Errorf("Could not get gs://%s/%s: %s", bucket, gsPath, err) |
| } |
| defer func() { assert.NoError(t, arch.Close()) }() |
| |
| // Open the output |
| var r io.ReadCloser = arch |
| if strings.HasSuffix(gsPath, ".gz") { |
| r, err = gzip.NewReader(r) |
| if err != nil { |
| return fmt.Errorf("Could not read gzip file: %s", err) |
| } |
| } |
| |
| f, err := os.Create(targetPath) |
| if err != nil { |
| return fmt.Errorf("Could not create target path: %s", err) |
| } |
| defer func() { assert.NoError(t, f.Close()) }() |
| _, err = io.Copy(f, r) |
| return err |
| } |
| |
| // DownloadTestDataArchive downloads testfiles that are stored in |
| // a gz compressed tar archive and decompresses them into the provided |
| // target directory. |
| func DownloadTestDataArchive(t assert.TestingT, bucket, gsPath, targetDir string) error { |
| if !strings.HasSuffix(gsPath, ".tar.gz") { |
| return fmt.Errorf("Expected .tar.gz file. But got:%s", gsPath) |
| } |
| |
| if err := os.MkdirAll(targetDir, 0755); err != nil { |
| return err |
| } |
| |
| arch, err := getStorangeItem(bucket, gsPath) |
| if err != nil { |
| return fmt.Errorf("Could not get gs://%s/%s: %s", bucket, gsPath, err) |
| } |
| defer func() { assert.NoError(t, arch.Close()) }() |
| |
| // Open the output |
| r, err := gzip.NewReader(arch) |
| if err != nil { |
| return fmt.Errorf("Could not read gzip archive: %s", err) |
| } |
| tarReader := tar.NewReader(r) |
| |
| for { |
| hdr, err := tarReader.Next() |
| if err == io.EOF { |
| break |
| } |
| |
| if err != nil { |
| return fmt.Errorf("Problem reading from tar archive: %s", err) |
| } |
| |
| targetPath := filepath.Join(targetDir, hdr.Name) |
| |
| if hdr.Typeflag == tar.TypeDir { |
| if err := os.MkdirAll(targetPath, 0755); err != nil { |
| return fmt.Errorf("Could not make %s: %s", targetPath, err) |
| } |
| } else { |
| f, err := os.Create(targetPath) |
| if err != nil { |
| return fmt.Errorf("Could not create target file %s: %s", targetPath, err) |
| } |
| _, err = io.Copy(f, tarReader) |
| if err != nil { |
| return fmt.Errorf("Problem while copying: %s", err) |
| } |
| defer func() { assert.NoError(t, f.Close()) }() |
| } |
| } |
| |
| return nil |
| } |
| |
| // MemoryGCSClient is a struct used for testing. Instead of writing to GCS, it |
| // stores data in memory. Not thread-safe. |
| type MemoryGCSClient struct { |
| bucket string |
| data map[string][]byte |
| } |
| |
| // Return a MemoryGCSClient instance. |
| func NewMemoryGCSClient(bucket string) *MemoryGCSClient { |
| return &MemoryGCSClient{ |
| bucket: bucket, |
| data: map[string][]byte{}, |
| } |
| } |
| |
| // See documentationn for GCSClient interface. |
| func (c *MemoryGCSClient) FileReader(ctx context.Context, path string) (io.ReadCloser, error) { |
| contents, ok := c.data[path] |
| if !ok { |
| return nil, storage.ErrObjectNotExist |
| } |
| return ioutil.NopCloser(bytes.NewReader(contents)), nil |
| } |
| |
| // io.WriteCloser implementation used by MemoryGCSClient. |
| type memoryWriter struct { |
| buf *bytes.Buffer |
| client *MemoryGCSClient |
| path string |
| } |
| |
| // See documentation for io.Writer. |
| func (w *memoryWriter) Write(p []byte) (int, error) { |
| return w.buf.Write(p) |
| } |
| |
| // See documentation for io.Closer. |
| func (w *memoryWriter) Close() error { |
| w.client.data[w.path] = w.buf.Bytes() |
| return nil |
| } |
| |
| // See documentation for GCSClient interface. |
| func (c *MemoryGCSClient) FileWriter(ctx context.Context, path string, opts FileWriteOptions) io.WriteCloser { |
| return &memoryWriter{ |
| buf: bytes.NewBuffer(nil), |
| client: c, |
| path: path, |
| } |
| } |
| |
| // See documentation for GCSClient interface. |
| func (c *MemoryGCSClient) DoesFileExist(ctx context.Context, path string) (bool, error) { |
| _, err := c.FileReader(ctx, path) |
| if err != nil { |
| return false, err |
| } |
| return true, nil |
| } |
| |
| // See documentation for GCSClient interface. |
| func (c *MemoryGCSClient) GetFileContents(ctx context.Context, path string) ([]byte, error) { |
| r, err := c.FileReader(ctx, path) |
| if err != nil { |
| return nil, err |
| } |
| return ioutil.ReadAll(r) |
| } |
| |
| // See documentation for GCSClient interface. |
| func (c *MemoryGCSClient) SetFileContents(ctx context.Context, path string, opts FileWriteOptions, contents []byte) error { |
| w := c.FileWriter(ctx, path, opts) |
| defer util.Close(w) |
| _, err := w.Write(contents) |
| return err |
| } |
| |
| // See documentation for GCSClient interface. |
| func (c *MemoryGCSClient) AllFilesInDirectory(ctx context.Context, prefix string, callback func(item *storage.ObjectAttrs)) error { |
| for key, data := range c.data { |
| if strings.HasPrefix(key, prefix) { |
| item := &storage.ObjectAttrs{ |
| Bucket: c.bucket, |
| Name: key, |
| Size: int64(len(data)), |
| } |
| callback(item) |
| } |
| } |
| return nil |
| } |
| |
| // See documentation for GCSClient interface. |
| func (c *MemoryGCSClient) DeleteFile(ctx context.Context, path string) error { |
| delete(c.data, path) |
| return nil |
| } |
| |
| // See documentation for GCSClient interface. |
| func (c *MemoryGCSClient) Bucket() string { |
| return c.bucket |
| } |