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