Add type parmeters to bqfakes (#163)

Author: stephen-soltesz

cloud/bqfake/bqfake.go

@@ -19,6 +19,11 @@ import (
 	"google.golang.org/api/iterator"
 )
 
+var (
+	// ErrTypeAssertionFailed may be returned by the row iterator when the target type is incorrect.
+	ErrTypeAssertionFailed = errors.New("type assertion failed")
+)
+
 // Table implements part of the bqiface.Table interface required for basic testing
 // Other parts of the interface should be implemented as needed.
 type Table struct {
@@ -99,44 +104,44 @@ func (ds Dataset) Table(name string) bqiface.Table {
 }
 
 // ClientConfig contains configuration for injecting result and error values.
-type ClientConfig struct {
-	QueryConfig
+type ClientConfig[Row any] struct {
+	QueryConfig[Row]
 }
 
 // QueryConfig contains configuration for injecting query results and error values.
-type QueryConfig struct {
+type QueryConfig[Row any] struct {
 	ReadErr error
-	RowIteratorConfig
+	RowIteratorConfig[Row]
 }
 
 // RowIteratorConfig contains configuration for injecting row iteration results and error values.
-type RowIteratorConfig struct {
+type RowIteratorConfig[Row any] struct {
 	IterErr error
-	Rows    []map[string]bigquery.Value
+	Rows    []Row
 }
 
 // Query implements parts of bqiface.Query to allow some very basic
 // unit tests.
-type Query struct {
+type Query[Row any] struct {
 	bqiface.Query
-	config QueryConfig
+	config QueryConfig[Row]
 }
 
 // SetQueryConfig is used to set the ReadErr or RowIteratorConfig.
-func (q Query) SetQueryConfig(bqiface.QueryConfig) {
+func (q Query[Row]) SetQueryConfig(bqiface.QueryConfig) {
 	log.Println("SetQueryConfig not implemented")
 }
 
-func (q Query) Run(context.Context) (bqiface.Job, error) {
+func (q Query[Row]) Run(context.Context) (bqiface.Job, error) {
 	log.Println("Run not implemented")
 	return Job{}, nil
 }
 
-func (q Query) Read(context.Context) (bqiface.RowIterator, error) {
+func (q Query[Row]) Read(context.Context) (bqiface.RowIterator, error) {
 	if q.config.ReadErr != nil {
 		return nil, q.config.ReadErr
 	}
-	return &RowIterator{config: q.config.RowIteratorConfig}, nil
+	return &RowIterator[Row]{config: q.config.RowIteratorConfig}, nil
 }
 
 // Job implements parts of bqiface.Job to allow some very basic
@@ -150,13 +155,17 @@ func (j Job) Wait(context.Context) (*bigquery.JobStatus, error) {
 	return nil, nil
 }
 
-type RowIterator struct {
+type RowIterator[Row any] struct {
 	bqiface.RowIterator
-	config RowIteratorConfig
+	config RowIteratorConfig[Row]
 	index  int
 }
 
-func (r *RowIterator) Next(dst interface{}) error {
+func (r *RowIterator[Row]) TotalRows() uint64 {
+	return uint64(len(r.config.Rows))
+}
+
+func (r *RowIterator[Row]) Next(dst interface{}) error {
 	// Check config for an error.
 	if r.config.IterErr != nil {
 		return r.config.IterErr
@@ -165,7 +174,10 @@ func (r *RowIterator) Next(dst interface{}) error {
 	if r.index >= len(r.config.Rows) {
 		return iterator.Done
 	}
-	v := dst.(*map[string]bigquery.Value)
+	v, ok := dst.(*Row)
+	if !ok {
+		return ErrTypeAssertionFailed
+	}
 	*v = r.config.Rows[r.index]
 	r.index++
 	return nil

cloud/bqfake/bqfake_test.go

@@ -29,7 +29,7 @@ func assertDataset(ds bqfake.Dataset) {
 }
 
 // This fails to compile if Client does not satisfy the interface.
-func assertClient(c bqfake.Client) {
+func assertClient[Row any](c bqfake.Client[Row]) {
 	func(cc bqiface.Client) {}(c)
 }
 
@@ -210,30 +210,32 @@ func TestQuery(t *testing.T) {
 func TestNewQueryReadClient(t *testing.T) {
 	tests := []struct {
 		name    string
-		config  bqfake.QueryConfig
+		config  bqfake.QueryConfig[map[string]bigquery.Value]
+		rows    uint64
 		wantErr bool
 	}{
 		{
 			name: "success",
-			config: bqfake.QueryConfig{
-				RowIteratorConfig: bqfake.RowIteratorConfig{
+			config: bqfake.QueryConfig[map[string]bigquery.Value]{
+				RowIteratorConfig: bqfake.RowIteratorConfig[map[string]bigquery.Value]{
 					Rows: []map[string]bigquery.Value{
 						{"okay": 1.234},
 					},
 				},
 			},
+			rows: 1,
 		},
 		{
 			name: "read-error",
-			config: bqfake.QueryConfig{
+			config: bqfake.QueryConfig[map[string]bigquery.Value]{
 				ReadErr: fmt.Errorf("Fake read error"),
 			},
 			wantErr: true,
 		},
 		{
 			name: "iter-error",
-			config: bqfake.QueryConfig{
-				RowIteratorConfig: bqfake.RowIteratorConfig{
+			config: bqfake.QueryConfig[map[string]bigquery.Value]{
+				RowIteratorConfig: bqfake.RowIteratorConfig[map[string]bigquery.Value]{
 					IterErr: fmt.Errorf("Fake iter error"),
 				},
 			},
@@ -252,6 +254,15 @@ func TestNewQueryReadClient(t *testing.T) {
 			if it == nil {
 				return
 			}
+			if it.TotalRows() != tt.rows {
+				t.Errorf("RowIterator.TotalRows() = %d, want %d", it.TotalRows(), tt.rows)
+			}
+			// Try loading wrong type.
+			var x string
+			err = it.Next(&x)
+			if !tt.wantErr && err != bqfake.ErrTypeAssertionFailed {
+				t.Errorf("RowIterator.Next() = %v, want %v", err, bqfake.ErrTypeAssertionFailed)
+			}
 			var row map[string]bigquery.Value
 			i := 0
 			for err = it.Next(&row); err == nil; err = it.Next(&row) {
@@ -265,7 +276,6 @@ func TestNewQueryReadClient(t *testing.T) {
 			if err == iterator.Done {
 				return
 			}
-			// err != nil.
 			if err != tt.config.RowIteratorConfig.IterErr {
 				t.Errorf("Next() error; got %v, want %v", err, tt.config.RowIteratorConfig.IterErr)
 				return

cloud/bqfake/client.go

@@ -66,40 +66,40 @@ func DryRunClient() (*http.Client, *CountingTransport) {
 }
 
 // Client implements a fake client.
-type Client struct {
+type Client[Row any] struct {
 	bqiface.Client
 	ctx    context.Context // Just for checking expiration/cancelation
-	config ClientConfig
+	config ClientConfig[Row]
 }
 
 // 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, opts ...option.ClientOption) (*Client[map[string]bigquery.Value], 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[map[string]bigquery.Value]{Client: bqiface.AdaptClient(c), ctx: ctx}, 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 {
+func (client Client[Row]) Dataset(ds string) bqiface.Dataset {
 	return Dataset{Dataset: client.Client.Dataset(ds), tables: make(map[string]*Table)}
 }
 
-func (client Client) Query(string) bqiface.Query {
-	return Query{
+func (client Client[Row]) Query(string) bqiface.Query {
+	return Query[Row]{
 		config: client.config.QueryConfig,
 	}
 }
 
-func NewQueryReadClient(qc QueryConfig) *Client {
+func NewQueryReadClient[Row any](qc QueryConfig[Row]) *Client[Row] {
 	// NOTE: if all needed functions are implemented by the fake, then a real
 	// client is unnecessary.
-	return &Client{
-		config: ClientConfig{QueryConfig: qc},
+	return &Client[Row]{
+		config: ClientConfig[Row]{QueryConfig: qc},
 	}
 }