Skip to content

Commit a97c930

Browse files
author
Jon Wayne Parrott
committed
Add initial simple test for uploading a file to a table
1 parent 4494e10 commit a97c930

File tree

1 file changed

+97
-0
lines changed

1 file changed

+97
-0
lines changed
Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
# Copyright 2017 Google Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import csv
16+
import io
17+
import json
18+
19+
from google import resumable_media
20+
import google.auth.transport.requests
21+
import mock
22+
import requests
23+
from six.moves import http_client
24+
from six.moves import StringIO
25+
26+
FIELDS = (u'name', u'age')
27+
ROWS = [
28+
(u'Phred Phlyntstone', 32),
29+
(u'Bharney Rhubble', 33),
30+
(u'Wylma Phlyntstone', 29),
31+
(u'Bhettye Rhubble', 27),
32+
]
33+
34+
35+
def rows_to_csv(fields, rows):
36+
"""Convert the rows into a CSV-format unicode string."""
37+
out = StringIO()
38+
writer = csv.writer(out)
39+
writer.writerow(fields)
40+
writer.writerows(rows)
41+
return out.getvalue()
42+
43+
44+
def make_table():
45+
from google.cloud.bigquery import _http
46+
from google.cloud.bigquery import client
47+
from google.cloud.bigquery import dataset
48+
from google.cloud.bigquery import table
49+
50+
mock_connection = mock.Mock(spec=_http.Connection)
51+
mock_client = mock.Mock(spec=client.Client)
52+
mock_client._connection = mock_connection
53+
mock_client._credentials = mock.sentinel.credentials
54+
mock_client.project = 'project_id'
55+
56+
dataset = dataset.Dataset('test_dataset', mock_client)
57+
table = table.Table('test_table', dataset)
58+
59+
return table
60+
61+
62+
def make_response(status_code, content, headers={}):
63+
return mock.Mock(
64+
content=content, headers=headers, status_code=status_code,
65+
spec=requests.Response)
66+
67+
68+
def make_resumable_upload_responses(size):
69+
resumable_url = 'http://test.invalid?upload_id=and-then-there-was-1'
70+
initial_response = make_response(
71+
http_client.OK, b'', {'location': resumable_url})
72+
data_response = make_response(
73+
resumable_media.PERMANENT_REDIRECT,
74+
b'', {'range': 'bytes=0-{:d}'.format(size - 1)})
75+
final_response = make_response(
76+
http_client.OK,
77+
json.dumps({'size': size}),
78+
{'Content-Type': 'application/json'})
79+
return [initial_response, data_response, final_response]
80+
81+
82+
def test_upload_from_file_simple():
83+
table = make_table()
84+
85+
csv_file = io.BytesIO(
86+
rows_to_csv(FIELDS, ROWS).encode('utf-8'))
87+
csv_file_size = len(csv_file.getvalue())
88+
89+
mock_transport = mock.Mock(
90+
spec=google.auth.transport.requests.AuthorizedSession)
91+
transport_patch = mock.patch.object(
92+
table, '_make_transport', return_value=mock_transport)
93+
94+
with transport_patch:
95+
mock_transport.request.side_effect = make_resumable_upload_responses(
96+
csv_file_size)
97+
table.upload_from_file(csv_file, source_format='CSV')

0 commit comments

Comments
 (0)