|
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235 |
- // Copyright 2015 Google LLC
- //
- // Licensed under the Apache License, Version 2.0 (the "License");
- // you may not use this file except in compliance with the License.
- // You may obtain a copy of the License at
- //
- // http://www.apache.org/licenses/LICENSE-2.0
- //
- // Unless required by applicable law or agreed to in writing, software
- // distributed under the License is distributed on an "AS IS" BASIS,
- // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- // See the License for the specific language governing permissions and
- // limitations under the License.
-
- package bigquery
-
- import (
- "errors"
- "testing"
-
- "github.com/google/go-cmp/cmp"
-
- "cloud.google.com/go/internal/testutil"
-
- "golang.org/x/net/context"
- bq "google.golang.org/api/bigquery/v2"
- "google.golang.org/api/iterator"
- )
-
- type pageFetcherArgs struct {
- table *Table
- schema Schema
- startIndex uint64
- pageSize int64
- pageToken string
- }
-
- // pageFetcherReadStub services read requests by returning data from an in-memory list of values.
- type pageFetcherReadStub struct {
- // values and pageTokens are used as sources of data to return in response to calls to readTabledata or readQuery.
- values [][][]Value // contains pages / rows / columns.
- pageTokens map[string]string // maps incoming page token to returned page token.
-
- // arguments are recorded for later inspection.
- calls []pageFetcherArgs
- }
-
- func (s *pageFetcherReadStub) fetchPage(ctx context.Context, t *Table, schema Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error) {
- s.calls = append(s.calls,
- pageFetcherArgs{t, schema, startIndex, pageSize, pageToken})
- result := &fetchPageResult{
- pageToken: s.pageTokens[pageToken],
- rows: s.values[0],
- }
- s.values = s.values[1:]
- return result, nil
- }
-
- func waitForQueryStub(context.Context, string) (Schema, error) {
- return nil, nil
- }
-
- func TestRead(t *testing.T) {
- // The data for the service stub to return is populated for each test case in the testCases for loop.
- ctx := context.Background()
- c := &Client{projectID: "project-id"}
- pf := &pageFetcherReadStub{}
- queryJob := &Job{
- projectID: "project-id",
- jobID: "job-id",
- c: c,
- config: &bq.JobConfiguration{
- Query: &bq.JobConfigurationQuery{
- DestinationTable: &bq.TableReference{
- ProjectId: "project-id",
- DatasetId: "dataset-id",
- TableId: "table-id",
- },
- },
- },
- }
-
- for _, readFunc := range []func() *RowIterator{
- func() *RowIterator {
- return c.Dataset("dataset-id").Table("table-id").read(ctx, pf.fetchPage)
- },
- func() *RowIterator {
- it, err := queryJob.read(ctx, waitForQueryStub, pf.fetchPage)
- if err != nil {
- t.Fatal(err)
- }
- return it
- },
- } {
- testCases := []struct {
- data [][][]Value
- pageTokens map[string]string
- want [][]Value
- }{
- {
- data: [][][]Value{{{1, 2}, {11, 12}}, {{30, 40}, {31, 41}}},
- pageTokens: map[string]string{"": "a", "a": ""},
- want: [][]Value{{1, 2}, {11, 12}, {30, 40}, {31, 41}},
- },
- {
- data: [][][]Value{{{1, 2}, {11, 12}}, {{30, 40}, {31, 41}}},
- pageTokens: map[string]string{"": ""}, // no more pages after first one.
- want: [][]Value{{1, 2}, {11, 12}},
- },
- }
- for _, tc := range testCases {
- pf.values = tc.data
- pf.pageTokens = tc.pageTokens
- if got, ok := collectValues(t, readFunc()); ok {
- if !testutil.Equal(got, tc.want) {
- t.Errorf("reading: got:\n%v\nwant:\n%v", got, tc.want)
- }
- }
- }
- }
- }
-
- func collectValues(t *testing.T, it *RowIterator) ([][]Value, bool) {
- var got [][]Value
- for {
- var vals []Value
- err := it.Next(&vals)
- if err == iterator.Done {
- break
- }
- if err != nil {
- t.Errorf("err calling Next: %v", err)
- return nil, false
- }
- got = append(got, vals)
- }
- return got, true
- }
-
- func TestNoMoreValues(t *testing.T) {
- c := &Client{projectID: "project-id"}
- pf := &pageFetcherReadStub{
- values: [][][]Value{{{1, 2}, {11, 12}}},
- }
- it := c.Dataset("dataset-id").Table("table-id").read(context.Background(), pf.fetchPage)
- var vals []Value
- // We expect to retrieve two values and then fail on the next attempt.
- if err := it.Next(&vals); err != nil {
- t.Fatalf("Next: got: %v: want: nil", err)
- }
- if err := it.Next(&vals); err != nil {
- t.Fatalf("Next: got: %v: want: nil", err)
- }
- if err := it.Next(&vals); err != iterator.Done {
- t.Fatalf("Next: got: %v: want: iterator.Done", err)
- }
- }
-
- var errBang = errors.New("bang!")
-
- func errorFetchPage(context.Context, *Table, Schema, uint64, int64, string) (*fetchPageResult, error) {
- return nil, errBang
- }
-
- func TestReadError(t *testing.T) {
- // test that service read errors are propagated back to the caller.
- c := &Client{projectID: "project-id"}
- it := c.Dataset("dataset-id").Table("table-id").read(context.Background(), errorFetchPage)
- var vals []Value
- if err := it.Next(&vals); err != errBang {
- t.Fatalf("Get: got: %v: want: %v", err, errBang)
- }
- }
-
- func TestReadTabledataOptions(t *testing.T) {
- // test that read options are propagated.
- s := &pageFetcherReadStub{
- values: [][][]Value{{{1, 2}}},
- }
- c := &Client{projectID: "project-id"}
- tr := c.Dataset("dataset-id").Table("table-id")
- it := tr.read(context.Background(), s.fetchPage)
- it.PageInfo().MaxSize = 5
- var vals []Value
- if err := it.Next(&vals); err != nil {
- t.Fatal(err)
- }
- want := []pageFetcherArgs{{
- table: tr,
- pageSize: 5,
- pageToken: "",
- }}
- if diff := testutil.Diff(s.calls, want, cmp.AllowUnexported(pageFetcherArgs{}, pageFetcherReadStub{}, Table{}, Client{})); diff != "" {
- t.Errorf("reading (got=-, want=+):\n%s", diff)
- }
- }
-
- func TestReadQueryOptions(t *testing.T) {
- // test that read options are propagated.
- c := &Client{projectID: "project-id"}
- pf := &pageFetcherReadStub{
- values: [][][]Value{{{1, 2}}},
- }
- tr := &bq.TableReference{
- ProjectId: "project-id",
- DatasetId: "dataset-id",
- TableId: "table-id",
- }
- queryJob := &Job{
- projectID: "project-id",
- jobID: "job-id",
- c: c,
- config: &bq.JobConfiguration{
- Query: &bq.JobConfigurationQuery{DestinationTable: tr},
- },
- }
- it, err := queryJob.read(context.Background(), waitForQueryStub, pf.fetchPage)
- if err != nil {
- t.Fatalf("err calling Read: %v", err)
- }
- it.PageInfo().MaxSize = 5
- var vals []Value
- if err := it.Next(&vals); err != nil {
- t.Fatalf("Next: got: %v: want: nil", err)
- }
-
- want := []pageFetcherArgs{{
- table: bqToTable(tr, c),
- pageSize: 5,
- pageToken: "",
- }}
- if !testutil.Equal(pf.calls, want, cmp.AllowUnexported(pageFetcherArgs{}, Table{}, Client{})) {
- t.Errorf("reading: got:\n%v\nwant:\n%v", pf.calls, want)
- }
- }
|