def test_stopping_criteria_with_startids(self, write_row):
yaml = """
- object: foo
just_once: True
- object: bar
"""
continuation_yaml = """
id_manager:
last_used_ids:
foo: 41
bar: 1000
nicknames_and_tables: {}
today: 2022-11-03
"""
generate(
StringIO(yaml),
continuation_file=StringIO(continuation_yaml),
stopping_criteria=StoppingCriteria("bar", 3),
)
assert write_row.mock_calls == [
mock.call("bar", {"id": 1001}),
mock.call("bar", {"id": 1002}),
mock.call("bar", {"id": 1003}),
]
def test_stopping_criteria_with_startids(self, write_row):
yaml = """
- object: foo
just_once: True
- object: bar
"""
continuation_yaml = """
!snowfakery_globals
id_manager: !snowfakery_ids
last_used_ids:
foo: 41
bar: 1000
last_seen_obj_of_type:
foo:
id: 41
bar:
id: 1000
nicknames_and_tables: {}
"""
generate(
StringIO(yaml),
continuation_file=StringIO(continuation_yaml),
stopping_criteria=StoppingCriteria("bar", 3),
)
assert write_row.mock_calls == [
mock.call("bar", {"id": 1001}),
mock.call("bar", {"id": 1002}),
mock.call("bar", {"id": 1003}),
]
def test_stops_on_no_progress(self, write_row):
yaml = """
- object: foo
just_once: True
- object: bar
"""
with self.assertRaises(RuntimeError):
generate(StringIO(yaml), stopping_criteria=StoppingCriteria("foo", 3))
def test_stops_if_criteria_misspelled(self, write_row):
yaml = """
- object: foo
just_once: True
- object: bar
"""
with self.assertRaises(DataGenError):
generate(StringIO(yaml),
stopping_criteria=StoppingCriteria("baz", 3))
def stopping_criteria_from_target_number(target_number):
"Deconstruct a tuple of 'str number' or 'number str' and make a StoppingCriteria"
# 'number str' is the official format so the other one can be deprecated one day.
if target_number:
if isinstance(target_number[0], int):
target_number = target_number[1], target_number[0]
return StoppingCriteria(*target_number)
return None
def _generate_loop(self, yaml, iterations, num_per_iteration):
old_continuation_data = None
for i in range(0, iterations):
continuation_file = StringIO()
generate(
StringIO(yaml),
{},
stopping_criteria=StoppingCriteria("A", num_per_iteration),
continuation_file=old_continuation_data,
generate_continuation_file=continuation_file,
)
old_continuation_data = StringIO(continuation_file.getvalue())
def test_stopping_criteria_and_just_once(self, write_row):
yaml = """
- object: foo
just_once: True
- object: bar
"""
generate(StringIO(yaml), stopping_criteria=StoppingCriteria("bar", 3))
assert write_row.mock_calls == [
mock.call("foo", {"id": 1}),
mock.call("bar", {"id": 1}),
mock.call("bar", {"id": 2}),
mock.call("bar", {"id": 3}),
]
def test_random_reference_local(self, generated_rows):
yaml = """ #1
- object: A #2
count: 10 #4
- object: B #5
fields: #7
A_ref: #8
random_reference: #9
tablename: A #10
scope: current-iteration #11
"""
with mock.patch("random.randint") as randint:
randint.side_effect = [8, 12]
generate(StringIO(yaml), stopping_criteria=StoppingCriteria("B", 2))
assert randint.mock_calls == [mock.call(1, 10), mock.call(11, 20)]
assert generated_rows.table_values("B", 2, "A_ref") == "A(12)"
def test_random_reference_global(self, generated_rows):
# undocumented experimental feature!
yaml = """ #1
- object: A #2
count: 10 #4
- object: B #5
count: 2 #6
fields: #7
A_ref: #8
random_reference:
tablename: A #9
scope: prior-and-current-iterations
"""
with mock.patch("random.randint") as randint, mock.patch(
"warnings.warn"
) as warn:
randint.side_effect = [8, 3, 8, 3]
generate(StringIO(yaml), stopping_criteria=StoppingCriteria("B", 4))
assert generated_rows.row_values(22, "A_ref") == "A(8)"
assert generated_rows.row_values(23, "A_ref") == "A(3)"
assert "experimental" in str(warn.mock_calls)
def test_forward_reference__iterations(self, generated_rows):
yaml = """
- object: A
fields:
B_ref:
reference:
B
- object: B
fields:
A_ref:
reference:
A
"""
generate(StringIO(yaml), {}, stopping_criteria=StoppingCriteria("A", 3))
assert generated_rows.mock_calls == [
mock.call("A", {"id": 1, "B_ref": "B(1)"}),
mock.call("B", {"id": 1, "A_ref": "A(1)"}),
mock.call("A", {"id": 2, "B_ref": "B(2)"}),
mock.call("B", {"id": 2, "A_ref": "A(2)"}),
mock.call("A", {"id": 3, "B_ref": "B(3)"}),
mock.call("B", {"id": 3, "A_ref": "A(3)"}),
]
def _init_options(self, kwargs):
super()._init_options(kwargs)
self.yaml_file = os.path.abspath(self.options["generator_yaml"])
if not os.path.exists(self.yaml_file):
raise TaskOptionsError(f"Cannot find {self.yaml_file}")
if "vars" in self.options:
self.vars = process_list_of_pairs_dict_arg(self.options["vars"])
self.generate_mapping_file = self.options.get("generate_mapping_file")
if self.generate_mapping_file:
self.generate_mapping_file = os.path.abspath(
self.generate_mapping_file)
num_records = self.options.get("num_records")
if num_records is not None:
num_records_tablename = self.options.get("num_records_tablename")
if not num_records_tablename:
raise TaskOptionsError(
f"Cannot specify num_records without num_records_tablename."
)
self.stopping_criteria = StoppingCriteria(num_records_tablename,
num_records)
self.working_directory = self.options.get("working_directory")
def generate_cli(
yaml_file,
option=[],
dburls=[],
target_number=None,
mapping_file=None,
debug_internals=False,
generate_cci_mapping_file=None,
output_format=None,
output_files=None,
output_folder=None,
nickname_ids=None,
continuation_file=None,
generate_continuation_file=None,
):
"""
Generates records from a YAML file
\b
Records can go to:
* stdout (default)
* JSON file (--output_format=json --output-file=foo.json)
* diagram file (--output_format=png --output-file=foo.png)
* a database identified by --dburl (e.g. --dburl sqlite:////tmp/foo.db)
* or to a directory as a set of CSV files (--output-format=csv --output-folder=csvfiles)
Diagram output depends on the installation of pygraphviz ("pip install pygraphviz")
"""
output_files = list(output_files) if output_files else []
stopping_criteria = StoppingCriteria(
*target_number) if target_number else None
output_format = output_format.lower() if output_format else None
validate_options(
yaml_file,
option,
dburls,
mapping_file,
debug_internals,
generate_cci_mapping_file,
output_format,
output_files,
output_folder,
)
with configure_output_stream(dburls, mapping_file, output_format,
output_files, output_folder) as output_stream:
try:
with click.open_file(yaml_file) as f:
summary = generate(
open_yaml_file=f,
user_options=dict(option),
output_stream=output_stream,
stopping_criteria=stopping_criteria,
generate_continuation_file=generate_continuation_file,
continuation_file=continuation_file,
)
if debug_internals:
debuginfo = yaml.dump(summary.summarize_for_debugging(),
sort_keys=False)
sys.stderr.write(debuginfo)
if generate_cci_mapping_file:
with click.open_file(generate_cci_mapping_file, "w") as f:
yaml.safe_dump(mapping_from_recipe_templates(summary),
f,
sort_keys=False)
except DataGenError as e:
if debug_internals:
raise e
else:
click.echo("")
click.echo(
"An error occurred. If you would like to see a Python traceback, "
"use the --debug-internals option.")
raise click.ClickException(str(e)) from e
