def testBigQueryToExample(self, mock_client):
# Mock query result schema for _BigQueryConverter.
mock_client.return_value.query.return_value.result.return_value.schema = self._schema
with beam.Pipeline() as pipeline:
examples = (
pipeline | 'ToTFExample' >> executor._BigQueryToExample(
input_dict={},
exec_properties={
'_beam_pipeline_args': ['--project=' + _test_project],
},
split_pattern='SELECT i, i2, b, f, f2, s, s2 FROM `fake`'))
feature = {}
feature['i'] = tf.train.Feature(int64_list=tf.train.Int64List(value=[1]))
feature['i2'] = tf.train.Feature(
int64_list=tf.train.Int64List(value=[2, 3]))
feature['b'] = tf.train.Feature(int64_list=tf.train.Int64List(value=[1]))
feature['f'] = tf.train.Feature(
float_list=tf.train.FloatList(value=[2.0]))
feature['f2'] = tf.train.Feature(
float_list=tf.train.FloatList(value=[2.7, 3.8]))
feature['s'] = tf.train.Feature(
bytes_list=tf.train.BytesList(value=[tf.compat.as_bytes('abc')]))
feature['s2'] = tf.train.Feature(
bytes_list=tf.train.BytesList(
value=[tf.compat.as_bytes('abc'),
tf.compat.as_bytes('def')]))
example_proto = tf.train.Example(
features=tf.train.Features(feature=feature))
util.assert_that(examples, util.equal_to([example_proto]))
def testBigQueryToExample(self, mock_client):
# Mock query result schema for _BigQueryConverter.
mock_client.return_value.query.return_value.result.return_value.schema = self._schema
with beam.Pipeline() as pipeline:
examples = (pipeline | 'ToTFExample' >>
executor._BigQueryToExample({}, self._exec_properties))
feature = {}
feature['i'] = tf.train.Feature(int64_list=tf.train.Int64List(
value=[1]))
feature['f'] = tf.train.Feature(float_list=tf.train.FloatList(
value=[2.0]))
feature['s'] = tf.train.Feature(bytes_list=tf.train.BytesList(
value=[tf.compat.as_bytes('abc')]))
example_proto = tf.train.Example(features=tf.train.Features(
feature=feature))
util.assert_that(examples, util.equal_to([example_proto]))