go/gcs/test_gcsclient/memory_client.go

package test_gcsclient

import (
	"bytes"
	"compress/gzip"
	"context"
	"io"
	"io/ioutil"
	"strings"
	"sync"

	"cloud.google.com/go/storage"
	"go.skia.org/infra/go/gcs"
	"go.skia.org/infra/go/util"
)

// MemoryGCSClient is a struct used for testing. Instead of writing to GCS, it
// stores data in memory.
type MemoryGCSClient struct {
	bucket string
	data   map[string][]byte
	opts   map[string]gcs.FileWriteOptions
	mtx    sync.RWMutex
}

// Return a MemoryGCSClient instance.
func NewMemoryClient(bucket string) *MemoryGCSClient {
	return &MemoryGCSClient{
		bucket: bucket,
		data:   map[string][]byte{},
		opts:   map[string]gcs.FileWriteOptions{},
	}
}

// See documentationn for GCSClient interface.
func (c *MemoryGCSClient) FileReader(ctx context.Context, path string) (io.ReadCloser, error) {
	c.mtx.RLock()
	defer c.mtx.RUnlock()
	contents, ok := c.data[path]
	if !ok {
		return nil, storage.ErrObjectNotExist
	}
	rv := ioutil.NopCloser(bytes.NewReader(contents))
	// GCS automatically decodes gzip-encoded files. See
	// https://cloud.google.com/storage/docs/transcoding. We do the same here so that tests acurately
	// reflect what will happen when actually using GCS.
	if c.opts[path].ContentEncoding == "gzip" {
		var err error
		rv, err = gzip.NewReader(rv)
		if err != nil {
			return nil, err
		}
	}
	return rv, 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.mtx.Lock()
	defer w.client.mtx.Unlock()
	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 gcs.FileWriteOptions) io.WriteCloser {
	c.mtx.Lock()
	defer c.mtx.Unlock()
	c.opts[path] = opts
	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 {
		if err == storage.ErrObjectNotExist {
			return false, 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 gcs.FileWriteOptions, contents []byte) error {
	return gcs.WithWriteFile(c, ctx, path, opts, func(w io.Writer) error {
		_, err := w.Write(contents)
		return err
	})
}

// See documentation for GCSClient interface.
func (c *MemoryGCSClient) GetFileObjectAttrs(ctx context.Context, path string) (*storage.ObjectAttrs, error) {
	c.mtx.RLock()
	defer c.mtx.RUnlock()
	data, ok := c.data[path]
	if !ok {
		return nil, storage.ErrObjectNotExist
	}
	opts, ok := c.opts[path]
	if !ok {
		return nil, storage.ErrObjectNotExist
	}
	return &storage.ObjectAttrs{
		Bucket:             c.bucket,
		Name:               path,
		ContentType:        opts.ContentType,
		ContentLanguage:    opts.ContentLanguage,
		Size:               int64(len(data)),
		ContentEncoding:    opts.ContentEncoding,
		ContentDisposition: opts.ContentDisposition,
		Metadata:           util.CopyStringMap(opts.Metadata),
	}, nil
}

// See documentation for GCSClient interface.
func (c *MemoryGCSClient) AllFilesInDirectory(ctx context.Context, prefix string, callback func(item *storage.ObjectAttrs)) error {
	items := func() []*storage.ObjectAttrs {
		c.mtx.RLock()
		defer c.mtx.RUnlock()
		var items []*storage.ObjectAttrs
		for key, data := range c.data {
			if strings.HasPrefix(key, prefix) {
				opts := c.opts[key]
				items = append(items, &storage.ObjectAttrs{
					Bucket:             c.bucket,
					Name:               key,
					ContentType:        opts.ContentType,
					ContentLanguage:    opts.ContentLanguage,
					Size:               int64(len(data)),
					ContentEncoding:    opts.ContentEncoding,
					ContentDisposition: opts.ContentDisposition,
					Metadata:           util.CopyStringMap(opts.Metadata),
				})
			}
		}
		return items
	}()
	for _, item := range items {
		callback(item)
	}
	return nil
}

// See documentation for GCSClient interface.
func (c *MemoryGCSClient) DeleteFile(ctx context.Context, path string) error {
	c.mtx.Lock()
	defer c.mtx.Unlock()
	delete(c.data, path)
	delete(c.opts, path)
	return nil
}

// See documentation for GCSClient interface.
func (c *MemoryGCSClient) Bucket() string {
	return c.bucket
}

// make sure MemoryGCSClient implements the GCSClient interface
var _ gcs.GCSClient = (*MemoryGCSClient)(nil)