Extend bqfake Client, Dataset, and Table (#164)

Author: cristinaleonr

cloudtest/bqfake/bqfake.go

@@ -27,6 +27,17 @@ type Table struct {
 	name string
 	// NOTE: TableType is used to indicate if this is initialized
 	metadata *bigquery.TableMetadata
+	err      error
+}
+
+// NewTable returns a new instance of Table.
+func NewTable(ds Dataset, name string, md *bigquery.TableMetadata, err error) *Table {
+	return &Table{
+		ds:       ds,
+		name:     name,
+		metadata: md,
+		err:      err,
+	}
 }
 
 // ProjectID implements the bqiface method.
@@ -80,10 +91,42 @@ func (tbl Table) Create(ctx context.Context, meta *bigquery.TableMetadata) error
 	return nil
 }
 
+// Update updates the table's `Schema`.
+func (tbl Table) Update(ctx context.Context, md bigquery.TableMetadataToUpdate, etag string) (*bigquery.TableMetadata, error) {
+	if md.Schema != nil {
+		tbl.metadata.Schema = md.Schema
+	}
+	return tbl.metadata, nil
+}
+
 // Dataset implements part of the bqiface.Dataset interface.
 type Dataset struct {
 	bqiface.Dataset
-	tables map[string]*Table
+	tables   map[string]*Table
+	metadata *bqiface.DatasetMetadata
+	err      error
+}
+
+// NewDataset returns a new instance of Dataset.
+func NewDataset(t map[string]*Table, md *bqiface.DatasetMetadata, err error) *Dataset {
+	return &Dataset{
+		tables:   t,
+		metadata: md,
+		err:      err,
+	}
+}
+
+// Metadata implements the bqiface method.
+func (ds Dataset) Metadata(ctx context.Context) (*bqiface.DatasetMetadata, error) {
+	if ds.metadata != nil {
+		return ds.metadata, nil
+	}
+	return nil, errors.New("invalid dataset metadata")
+}
+
+// Create returns an error.
+func (ds Dataset) Create(ctx context.Context, md *bqiface.DatasetMetadata) error {
+	return ds.err
 }
 
 // Table implements the bqiface method.

cloudtest/bqfake/bqfake_test.go

@@ -2,17 +2,19 @@ package bqfake_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"net/http"
 	"reflect"
 	"testing"
 	"time"
 
-	"google.golang.org/api/option"
 	"cloud.google.com/go/bigquery"
+	"github.com/google/go-cmp/cmp"
 	"github.com/googleapis/google-cloud-go-testing/bigquery/bqiface"
 	"google.golang.org/api/iterator"
+	"google.golang.org/api/option"
 
 	"github.com/m-lab/go/cloudtest/bqfake"
 )
@@ -55,7 +57,7 @@ func TestNewClientErr(t *testing.T) {
 	ctx := context.Background()
 	// These options are incompatible with one another and generate an error from bigquery.NewClient.
 	opts := []option.ClientOption{option.WithAPIKey("asdf"), option.WithoutAuthentication()}
-	c, err := bqfake.NewClient(ctx, "fakeProject", opts...)
+	c, err := bqfake.NewClient(ctx, "fakeProject", map[string]*bqfake.Dataset{}, opts...)
 	if err == nil {
 		c.Close()
 		t.Fatal("Should return constructing client error")
@@ -74,9 +76,82 @@ func TestBadDataset(t *testing.T) {
 	ds.Table("Foobar")
 }
 
+func TestDatasetMetadata(t *testing.T) {
+	tests := []struct {
+		name    string
+		md      *bqiface.DatasetMetadata
+		want    *bqiface.DatasetMetadata
+		wantErr bool
+	}{
+		{
+			name: "success",
+			md: &bqiface.DatasetMetadata{
+				DatasetMetadata: bigquery.DatasetMetadata{
+					Name: "fakeMetadata",
+				},
+			},
+			want: &bqiface.DatasetMetadata{
+				DatasetMetadata: bigquery.DatasetMetadata{
+					Name: "fakeMetadata",
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name:    "error",
+			md:      nil,
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ds := bqfake.NewDataset(nil, tt.md, nil)
+			got, err := ds.Metadata(context.Background())
+			if (err != nil) != tt.wantErr {
+				t.Errorf("Dataset.Metadata() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+
+			if !tt.wantErr && !cmp.Equal(got, tt.want) {
+				t.Errorf("Dataset.Metadata() got = %v, want = %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestDatasetCreate(t *testing.T) {
+	tests := []struct {
+		name    string
+		err     error
+		wantErr bool
+	}{
+		{
+			name:    "success",
+			err:     nil,
+			wantErr: false,
+		},
+		{
+			name:    "error",
+			err:     errors.New("fakeError"),
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ds := bqfake.NewDataset(nil, nil, tt.err)
+			got := ds.Create(context.Background(), &bqiface.DatasetMetadata{})
+			if (got != nil) != tt.wantErr {
+				t.Errorf("Dataset.Create() error = %v, wantErr %v", got, tt.wantErr)
+			}
+		})
+	}
+}
+
 func TestUninitializedTable(t *testing.T) {
 	ctx := context.Background()
-	c, err := bqfake.NewClient(ctx, "fakeProject")
+	c, err := bqfake.NewClient(ctx, "fakeProject", map[string]*bqfake.Dataset{})
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -113,7 +188,7 @@ func TestUninitializedTable(t *testing.T) {
 
 func TestTable(t *testing.T) {
 	ctx := context.Background()
-	c, err := bqfake.NewClient(ctx, "fakeProject")
+	c, err := bqfake.NewClient(ctx, "fakeProject", map[string]*bqfake.Dataset{})
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -149,7 +224,7 @@ func createTable(ctx context.Context, ds bqiface.Dataset, name string) error {
 
 func TestTableMetadata(t *testing.T) {
 	ctx := context.Background()
-	c, err := bqfake.NewClient(ctx, "fakeProject")
+	c, err := bqfake.NewClient(ctx, "fakeProject", map[string]*bqfake.Dataset{})
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -179,9 +254,31 @@ func TestTableMetadata(t *testing.T) {
 	}
 }
 
+func TestTableUpdate(t *testing.T) {
+	want := &bigquery.TableMetadata{
+		Schema: []*bigquery.FieldSchema{
+			{
+				Name: "fakeField",
+			},
+		},
+	}
+	tbl := bqfake.NewTable(bqfake.Dataset{}, "", &bigquery.TableMetadata{}, nil)
+	got, err := tbl.Update(context.Background(), bigquery.TableMetadataToUpdate{
+		Schema: want.Schema,
+	}, "")
+
+	if err != nil {
+		t.Errorf("Table.Update() error = %v, wantErr = nil", err)
+	}
+
+	if !cmp.Equal(got, want) {
+		t.Errorf("Table.Update() got = %v, want = %v", got, want)
+	}
+}
+
 func TestQuery(t *testing.T) {
 	ctx := context.Background()
-	c, err := bqfake.NewClient(ctx, "fakeProject")
+	c, err := bqfake.NewClient(ctx, "fakeProject", map[string]*bqfake.Dataset{})
 	if err != nil {
 		t.Fatal(err)
 	}

cloudtest/bqfake/client.go

@@ -68,26 +68,32 @@ func DryRunClient() (*http.Client, *CountingTransport) {
 // Client implements a fake client.
 type Client struct {
 	bqiface.Client
-	ctx    context.Context // Just for checking expiration/cancelation
-	config ClientConfig
+	ctx      context.Context // Just for checking expiration/cancelation
+	datasets map[string]*Dataset
+	config   ClientConfig
 }
 
 // NewClient creates a new Client implementing bqiface.Client, with a dry run HTTPClient.
-func NewClient(ctx context.Context, project string, opts ...option.ClientOption) (*Client, error) {
+func NewClient(ctx context.Context, project string, ds map[string]*Dataset, opts ...option.ClientOption) (*Client, error) {
 	dryRun, _ := DryRunClient()
 	opts = append(opts, option.WithHTTPClient(dryRun))
 	c, err := bigquery.NewClient(ctx, project, opts...)
 	if err != nil {
 		return nil, err
 	}
-	return &Client{Client: bqiface.AdaptClient(c), ctx: ctx}, nil
+	return &Client{Client: bqiface.AdaptClient(c), ctx: ctx, datasets: ds}, nil
 }
 
 // Dataset creates a Dataset.
 // TODO - understand how bqiface adapters/structs work, and make this return a Dataset
 // that satisfies bqiface.Dataset interface?
-func (client Client) Dataset(ds string) bqiface.Dataset {
-	return Dataset{Dataset: client.Client.Dataset(ds), tables: make(map[string]*Table)}
+func (client Client) Dataset(name string) bqiface.Dataset {
+	ds, ok := client.datasets[name]
+	if !ok {
+		ds = &Dataset{Dataset: client.Client.Dataset(name), tables: make(map[string]*Table)}
+		client.datasets[name] = ds
+	}
+	return ds
 }
 
 func (client Client) Query(string) bqiface.Query {