def test_on_direct_runner(self):
class FakeSink(NativeSink):
"""A fake sink outputing a number of elements."""
def __init__(self):
self.written_values = []
self.writer_instance = FakeSinkWriter(self.written_values)
def writer(self):
return self.writer_instance
class FakeSinkWriter(NativeSinkWriter):
"""A fake sink writer for testing."""
def __init__(self, written_values):
self.written_values = written_values
def __enter__(self):
return self
def __exit__(self, *unused_args):
pass
def Write(self, value):
self.written_values.append(value)
with TestPipeline() as p:
sink = FakeSink()
p | Create(['a', 'b', 'c']) | _NativeWrite(sink) # pylint: disable=expression-not-assigned
self.assertEqual(['a', 'b', 'c'], sorted(sink.written_values))
def test_on_direct_runner(self):
class FakeSink(NativeSink):
"""A fake sink outputing a number of elements."""
def __init__(self):
self.written_values = []
self.writer_instance = FakeSinkWriter(self.written_values)
def writer(self):
return self.writer_instance
class FakeSinkWriter(NativeSinkWriter):
"""A fake sink writer for testing."""
def __init__(self, written_values):
self.written_values = written_values
def __enter__(self):
return self
def __exit__(self, *unused_args):
pass
def Write(self, value):
self.written_values.append(value)
p = TestPipeline()
sink = FakeSink()
p | Create(['a', 'b', 'c']) | _NativeWrite(sink) # pylint: disable=expression-not-assigned
p.run()
self.assertEqual(['a', 'b', 'c'], sink.written_values)
def test_on_direct_runner(self):
class FakeSink(NativeSink):
"""A fake sink outputing a number of elements."""
def __init__(self):
self.written_values = []
self.writer_instance = FakeSinkWriter(self.written_values)
def writer(self):
return self.writer_instance
class FakeSinkWriter(NativeSinkWriter):
"""A fake sink writer for testing."""
def __init__(self, written_values):
self.written_values = written_values
def __enter__(self):
return self
def __exit__(self, *unused_args):
pass
def Write(self, value):
self.written_values.append(value)
# Records in-memory writes, only works on Direct runner.
p = TestPipeline(runner='DirectRunner')
sink = FakeSink()
p | Create(['a', 'b', 'c']) | _NativeWrite(sink) # pylint: disable=expression-not-assigned
p.run()
self.assertEqual(['a', 'b', 'c'], sink.written_values)
def expand(self, pcoll):
from apache_beam.runners.dataflow.native_io import iobase as dataflow_io
if isinstance(self.sink, dataflow_io.NativeSink):
# A native sink
return pcoll | 'NativeWrite' >> dataflow_io._NativeWrite(self.sink)
elif isinstance(self.sink, Sink):
# A custom sink
return pcoll | WriteImpl(self.sink)
elif isinstance(self.sink, ptransform.PTransform):
# This allows "composite" sinks to be used like non-composite ones.
return pcoll | self.sink
else:
raise ValueError('A sink must inherit iobase.Sink, iobase.NativeSink, '
'or be a PTransform. Received : %r', self.sink)
def expand(self, pcoll):
from apache_beam.runners.dataflow.native_io import iobase as dataflow_io
if isinstance(self.sink, dataflow_io.NativeSink):
# A native sink
return pcoll | 'NativeWrite' >> dataflow_io._NativeWrite(self.sink)
elif isinstance(self.sink, Sink):
# A custom sink
return pcoll | WriteImpl(self.sink)
elif isinstance(self.sink, ptransform.PTransform):
# This allows "composite" sinks to be used like non-composite ones.
return pcoll | self.sink
else:
raise ValueError('A sink must inherit iobase.Sink, iobase.NativeSink, '
'or be a PTransform. Received : %r' % self.sink)
def setUp(self):
self.native_sink = NativeSink()
self.native_write = _NativeWrite(self.native_sink)
def setUp(self): self.native_sink = NativeSink() self.native_write = _NativeWrite(self.native_sink)
