def test_json_output_mocked(self):
yaml = """
- object: foo
count: 2
fields:
a: b
c: 3
"""
stdout = StringIO()
output_stream = JSONOutputStream(stdout)
generate(StringIO(yaml), {}, output_stream)
output_stream.close()
assert json.loads(stdout.getvalue()) == [
{
"_table": "foo",
"a": "b",
"c": 3.0,
"id": 1
},
{
"_table": "foo",
"a": "b",
"c": 3.0,
"id": 2
},
]
def configure_output_stream(dburls, mapping_file, output_format, output_files,
output_folder):
assert isinstance(output_files, (list, type(None)))
output_streams = [] # we allow multiple output streams
for dburl in dburls:
if mapping_file:
with click.open_file(mapping_file, "r") as f:
mappings = yaml.safe_load(f)
else:
mappings = None
output_streams.append(SqlOutputStream.from_url(dburl, mappings))
# JSON is the only output format (other than debug) that can go on stdout
if output_format == "json" and not output_files:
output_streams.append(JSONOutputStream(sys.stdout))
if output_format == "csv":
output_streams.append(CSVOutputStream(output_folder))
if output_files:
for path in output_files:
if output_folder:
path = Path(output_folder,
path) # put the file in the output folder
format = output_format or Path(path).suffix[1:]
if format == "json":
output_streams.append(JSONOutputStream(path))
elif format == "txt":
output_streams.append(DebugOutputStream(path))
elif format == "dot":
output_streams.append(GraphvizOutputStream(path))
elif format in graphic_file_extensions:
output_streams.append(ImageOutputStream(path, format))
else:
raise click.ClickException(
f"Unknown format or file extension: {format}")
if len(output_streams) == 0:
output_stream = DebugOutputStream()
elif len(output_streams) == 1:
output_stream = output_streams[0]
else:
output_stream = MultiplexOutputStream(output_streams)
try:
yield output_stream
finally:
for output_stream in output_streams:
try:
messages = output_stream.close()
except Exception as e:
messages = None
click.echo(f"Could not close {output_stream}: {str(e)}",
err=True)
if messages:
for message in messages:
click.echo(message)
def test_error_generates_empty_string(self):
yaml = """
- object: foo
fields: bar: baz: jaz
is_null:
"""
with StringIO() as s:
output_stream = JSONOutputStream(s)
with pytest.raises(exc.DataGenYamlSyntaxError):
generate(StringIO(yaml), {}, output_stream)
output_stream.close()
assert s.getvalue() == ""
def test_json_output_real(self):
yaml = """
- object: foo
count: 15
fields:
a: b
c: 3
"""
output_stream = JSONOutputStream(StringIO())
generate(StringIO(yaml), {}, output_stream)
output_stream.close()
def test_stringification(self, write_row):
yaml = """
- plugin: tests.test_custom_plugins_and_providers.EvalPlugin
- object: OBJ
fields:
some_value:
- EvalPlugin.add:
- 1
- EvalPlugin.sub:
- 5
- 3
"""
with StringIO() as s:
output_stream = JSONOutputStream(s)
generate(StringIO(yaml), {}, output_stream)
output_stream.close()
assert eval(s.getvalue())[0]["some_value"] == 3
def do_output(self, yaml):
with StringIO() as s:
output_stream = JSONOutputStream(s)
results = generate(StringIO(yaml), {}, output_stream)
output_stream.close()
return JSONTables(s.getvalue(), results.tables.keys())
