All tests pass.

Author: Tim Band

datafaker/create.py

@@ -9,16 +9,15 @@
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.orm import Session
 from sqlalchemy.schema import CreateColumn, CreateSchema, CreateTable, MetaData, Table
+import typer
 
 from datafaker.base import FileUploader
-from datafaker.make import get_generation_info
+from datafaker.make import get_generation_info, StoryGeneratorInfo
 from datafaker.populate import (
-    StoryGeneratorInfo,
     TableGenerator,
     call_function,
     get_symbols,
     get_table_generator_dict,
-    get_vocab_dict,
 )
 from datafaker.settings import get_destination_dsn, get_destination_schema, get_settings
 from datafaker.utils import (
@@ -157,8 +156,15 @@ def create_db_data(
 ) -> RowCounts:
     """Connect to a database and populate it with data."""
     if src_stats_filename:
-        with src_stats_filename.open(encoding="utf-8") as fh:
-            src_stats = yaml.load(fh, yaml.SafeLoader)
+        try:
+            with src_stats_filename.open(encoding="utf-8") as fh:
+                src_stats = yaml.load(fh, yaml.SafeLoader)
+        except FileNotFoundError:
+            logger.error(
+                "No source stats file '%', this should be the output of the 'make-stats' command",
+                src_stats_filename,
+            )
+            raise typer.Exit(1)
     else:
         src_stats = None
     return create_db_data_into(
@@ -187,10 +193,8 @@ def create_db_data_into(
 
     :param sorted_tables: The table names to populate, sorted so that foreign
         keys' targets are populated before the foreign keys themselves.
-    :param table_generator_dict: A mapping  of table names to the generators
-        used to make data for them.
-    :param story_generator_list: A list of story generators to be run after the
-        table generators on each pass.
+    :param config: The data from the ``config.yaml`` file.
+    :param src_stats: The data from the ``src-stats.yaml`` file.
     :param num_passes: Number of passes to perform.
     :param db_dsn: Connection string for the destination database.
     :param schema_name: Destination schema name.

datafaker/make.py

@@ -656,7 +656,7 @@ def get_generation_info(
     story_generators = _get_story_generators(config)
 
     max_unique_constraint_tries = get_property(
-        config, "max-unique-constraint-tries", str | None, None
+        config, "max-unique-constraint-tries", int | None, None
     )
     return GenerationInfo(
         provider_imports=PROVIDER_IMPORTS,

datafaker/populate.py

@@ -6,7 +6,7 @@
 from typing import Any, Callable
 
 from datafaker.base import FileUploader, ColumnPresence
-from datafaker.make import FunctionCall, TableGeneratorInfo, StoryGeneratorInfo
+from datafaker.make import FunctionCall, TableGeneratorInfo
 
 from datafaker.providers import (
     BytesProvider,
@@ -18,17 +18,34 @@
     TimespanProvider,
     WeightedBooleanProvider,
 )
-from datafaker.utils import logging, get_vocabulary_table_names, import_file
+from datafaker.utils import get_vocabulary_table_names, import_file
+
+def make_generic():
+    g = Generic(locale=Locale.EN_GB)
+    g.add_providers(
+        BytesProvider,
+        ColumnValueProvider,
+        DistributionProvider,
+        NullProvider,
+        SQLGroupByProvider,
+        TimedeltaProvider,
+        TimespanProvider,
+        WeightedBooleanProvider,
+    )
+    return g
+
+
+generic = make_generic()
 
-generic = Generic(locale=Locale.EN_GB)
 
-generic.add_provider(BytesProvider)
-generic.add_provider(ColumnValueProvider)
-generic.add_provider(NullProvider)
-generic.add_provider(SQLGroupByProvider)
-generic.add_provider(TimedeltaProvider)
-generic.add_provider(TimespanProvider)
-generic.add_provider(WeightedBooleanProvider)
+def reset_generic():
+    """
+    Reset all the generators.
+
+    Only really useful in test code.
+    """
+    global generic
+    generic = make_generic()
 
 
 def _eval_structure(config: Any, context: Mapping) -> Any:
@@ -57,10 +74,10 @@ def _eval_structure(config: Any, context: Mapping) -> Any:
 
 def _get_object(class_name: str, context: Mapping) -> Any:
     """
-    Get an object out of the context.
+    Fetch an object from the context.
 
     :param class_name: The name of the class, qualified if necessary.
-    Like "module.MyClass.Nested"
+      Like "module.MyClass.Nested"
     :param context: Mapping of strings to objects with those names.
     :return: A value from ``context`` if there are no qualifying names,
     otherwise the attribute of the base object.
@@ -288,11 +305,3 @@ def get_table_generator_dict(
         )
         for table_data in tables_data
     }
-
-
-def get_vocab_dict(config: Mapping, metadata: sqlalchemy.MetaData) -> Mapping[str, FileUploader]:
-    """Get a dict of table names to objects that can populate those tables from YAML files."""
-    return {
-        name: FileUploader(metadata.tables[name])
-        for name in get_vocabulary_table_names(config)
-    }

datafaker/utils.py

@@ -546,9 +546,7 @@ def get_row_generators(
     :param table_config: The element from the ``tables:`` stanza of ``config.xml``.
     :return: Pair of (name, row generator config).
     """
-    rgs = table_config.get("row_generators", None)
-    if isinstance(rgs, str) or not hasattr(rgs, "__iter__"):
-        return
+    rgs = get_property(table_config, "row_generators", list, [])
     for rg in rgs:
         name = rg.get("name", None)
         if name:

tests/test_create.py

@@ -21,6 +21,7 @@
     create_db_vocab,
     populate,
 )
+from datafaker.make import FunctionCall, StoryGeneratorInfo
 from datafaker.serialize_metadata import metadata_to_dict, dict_to_metadata
 from datafaker.utils import sorted_non_vocabulary_tables
 from tests.utils import DatafakerTestCase, GeneratesDBTestCase, RequiresDBTestCase
@@ -99,17 +100,16 @@ class TestPopulate(DatafakerTestCase):
     """Test create.populate."""
 
     # pylint: disable=too-many-locals
-    def test_populate(self) -> None:
+    @patch("datafaker.populate._get_object")
+    def test_populate(self, mock_get_object: MagicMock) -> None:
         """Test the populate function."""
         table_name = "table_name"
 
         def story() -> Generator[Tuple[str, dict], None, None]:
             """Mock story."""
             yield table_name, {}
 
-        def mock_story_gen(_: Any) -> Generator[Tuple[str, dict], None, None]:
-            """A function that returns mock stories."""
-            return story()
+        mock_get_object.return_value = story
 
         for num_stories_per_pass, num_rows_per_pass, num_initial_rows in itt.product(
             [0, 2], [0, 3], [0, 17]
@@ -130,11 +130,11 @@ def mock_story_gen(_: Any) -> Generator[Tuple[str, dict], None, None]:
 
                 story_generators: list[dict[str, Any]] = (
                     [
-                        {
-                            "function": mock_story_gen,
-                            "num_stories_per_pass": num_stories_per_pass,
-                            "name": "mock_story_gen",
-                        }
+                        StoryGeneratorInfo(
+                            "mock_story_gen name",
+                            FunctionCall("mock_story_gen", [], {}),
+                            num_stories_per_pass,
+                        )
                     ]
                     if num_stories_per_pass > 0
                     else []
@@ -304,11 +304,11 @@ def __call__(
         create_db_data_into(
             [MagicMock()],
             MagicMock(),
+            None,
             1,
             "duckdb:///:memory:data",
             None,
             MagicMock(),
-            MagicMock(),
         )
         assert mock_populate.side_effect.called
 
@@ -366,11 +366,12 @@ def test_create_data_minimal(self) -> None:
         }
         metadata = dict_to_metadata(orm, config)
         create_db_tables_into(metadata, self.dsn, self.schema_name)
+        generate_count = 4
         row_counts = create_db_data_into(
             sorted_non_vocabulary_tables(metadata, config),
             config,
             None,
-            4,
+            generate_count,
             self.dsn,
             self.schema_name,
             metadata,
@@ -379,3 +380,52 @@ def test_create_data_minimal(self) -> None:
             stmt = select(metadata.tables["one"])
             rows = connection.execute(stmt).fetchall()
             self.assertListEqual(rows, [(1,), (2,), (3,), (4,)])
+        self.assertListEqual(list(row_counts.keys()), ['one'])
+        self.assertEqual(row_counts["one"], generate_count)
+
+    def test_unique_constraint_minimal(self) -> None:
+        config = {
+            "tables": {
+                "one": {
+                    "row_generators": [{
+                        "name": "dist_gen.constant",
+                        "kwargs": {
+                            "value": 123,
+                        },
+                        "columns_assigned": ["tiger"],
+                    }]
+                }
+            },
+            "max-unique-constraint-tries": 20,
+        }
+        orm = {
+            "tables": {
+                "one": {
+                    "columns": {
+                        "id": {
+                            "primary": True,
+                            "type": "INTEGER",
+                        },
+                        "tiger": {
+                            "type": "INTEGER",
+                        },
+                    },
+                    "unique": [
+                        {"name": "tiger_uniq", "columns": ["tiger"]}
+                    ]
+                }
+            }
+        }
+        metadata = dict_to_metadata(orm, config)
+        create_db_tables_into(metadata, self.dsn, self.schema_name)
+        self.assertRaises(
+            RuntimeError,
+            create_db_data_into,
+            sorted_non_vocabulary_tables(metadata, config),
+            config,
+            None,
+            2,
+            self.dsn,
+            self.schema_name,
+            metadata,
+        )

tests/test_functional.py

@@ -272,52 +272,34 @@ def test_workflow_maximal_args(self) -> None:
             f"--config-file={self.config_file_path}",
             "--num-passes=2",
         )
-        self.assertEqual("", completed_process.stderr)
-        self.assertEqual(
-            sorted(
-                [
-                    "Creating data.",
-                    "Generating data for story 'story_generators.short_story'",
-                    "Generating data for story 'story_generators.short_story'",
-                    "Generating data for story 'story_generators.short_story'",
-                    "Generating data for story 'story_generators.short_story'",
-                    "Generating data for story 'story_generators.short_story'",
-                    "Generating data for story 'story_generators.short_story'",
-                    "Generating data for story 'story_generators.full_row_story'",
-                    "Generating data for story 'story_generators.full_row_story'",
-                    "Generating data for story 'story_generators.long_story'",
-                    "Generating data for story 'story_generators.long_story'",
-                    "Generating data for story 'story_generators.long_story'",
-                    "Generating data for story 'story_generators.long_story'",
-                    "Generating data for table 'data_type_test'",
-                    "Generating data for table 'data_type_test'",
-                    "Generating data for table 'no_pk_test'",
-                    "Generating data for table 'no_pk_test'",
-                    "Generating data for table 'person'",
-                    "Generating data for table 'person'",
-                    "Generating data for table 'strange_type_table'",
-                    "Generating data for table 'strange_type_table'",
-                    "Generating data for table 'unique_constraint_test'",
-                    "Generating data for table 'unique_constraint_test'",
-                    "Generating data for table 'unique_constraint_test2'",
-                    "Generating data for table 'unique_constraint_test2'",
-                    "Generating data for table 'test_entity'",
-                    "Generating data for table 'test_entity'",
-                    "Generating data for table 'hospital_visit'",
-                    "Generating data for table 'hospital_visit'",
-                    "Data created in 2 passes.",
-                    f"person: {2*(3+1+2+2)} rows created.",
-                    f"hospital_visit: {2*(2*2+3)} rows created.",
-                    "data_type_test: 2 rows created.",
-                    "no_pk_test: 2 rows created.",
-                    "strange_type_table: 2 rows created.",
-                    "unique_constraint_test: 2 rows created.",
-                    "unique_constraint_test2: 2 rows created.",
-                    "test_entity: 2 rows created.",
-                    "",
-                ]
-            ),
-            sorted(completed_process.stdout.split("\n")),
+        self.assertSetEqual(
+            {
+                "Creating data.",
+                "Generating data for story 'story_generators.short_story'",
+                "Generating data for story 'story_generators.full_row_story'",
+                "Generating data for story 'story_generators.full_row_story'",
+                "Generating data for story 'story_generators.long_story'",
+                "Generating data for table 'data_type_test'",
+                "Generating data for table 'no_pk_test'",
+                "Generating data for table 'person'",
+                "Generating data for table 'person'",
+                "Generating data for table 'strange_type_table'",
+                "Generating data for table 'unique_constraint_test'",
+                "Generating data for table 'unique_constraint_test2'",
+                "Generating data for table 'test_entity'",
+                "Generating data for table 'hospital_visit'",
+                "Data created in 2 passes.",
+                f"person: {2*(3+1+2+2)} rows created.",
+                f"hospital_visit: {2*(2*2+3)} rows created.",
+                "data_type_test: 2 rows created.",
+                "no_pk_test: 2 rows created.",
+                "strange_type_table: 2 rows created.",
+                "unique_constraint_test: 2 rows created.",
+                "unique_constraint_test2: 2 rows created.",
+                "test_entity: 2 rows created.",
+                "",
+            },
+            set(completed_process.stdout.split("\n")),
         )
 
         completed_process = self.invoke(
@@ -468,13 +450,9 @@ def test_unique_constraint_fail(self) -> None:
             f"--stats-file={self.stats_file_path}",
             "--num-passes=1",
         )
-        self.assertEqual("", completed_process.stderr)
         self.assertEqual(
-            "Generating data for story 'story_generators.short_story'\n"
-            "Generating data for story 'story_generators.short_story'\n"
             "Generating data for story 'story_generators.short_story'\n"
             "Generating data for story 'story_generators.full_row_story'\n"
-            "Generating data for story 'story_generators.long_story'\n"
             "Generating data for story 'story_generators.long_story'\n",
             completed_process.stdout,
         )
@@ -483,17 +461,14 @@ def test_unique_constraint_fail(self) -> None:
             "create-data",
             f"--config-file={self.config_file_path}",
             f"--orm-file={self.alt_orm_file_path}",
+            f"--stats-file={self.stats_file_path}",
             "--num-passes=3",
         )
-        self.assertEqual("", completed_process.stderr)
         self.assertEqual(
             (
-                "Generating data for story 'story_generators.short_story'\n"
-                "Generating data for story 'story_generators.short_story'\n"
                 "Generating data for story 'story_generators.short_story'\n"
                 "Generating data for story 'story_generators.full_row_story'\n"
                 "Generating data for story 'story_generators.long_story'\n"
-                "Generating data for story 'story_generators.long_story'\n"
             )
             * 3,
             completed_process.stdout,
@@ -504,6 +479,7 @@ def test_unique_constraint_fail(self) -> None:
             "create-data",
             f"--config-file={self.config_file_path}",
             f"--orm-file={self.alt_orm_file_path}",
+            f"--stats-file={self.stats_file_path}",
             "--num-passes=1",
             expected_error=(
                 "Failed to satisfy unique constraints for table unique_constraint_test"