Repository: SFDO-Tooling/CumulusCI
Pull request #3569 (2023-04-04 00:20:13)

The `capture_sample_data` task now looks for a `datasetname.load.yml` file to influence how it generates its .mapping file
CumulusCI has a feature called a ".[load.yml](https://cumulusci.readthedocs.io/en/stable/data.html#controlling-the-loading-process)" file which is used when generating mapping files in the context of loading Snowfakery recipes. This allows you to, for example, control the API used, change the load order, etc. Basically it allows you to override default loading behaviours.

Our new `capture_sample_data` task also generates mapping files, and therefore could benefit from the same kinds of overrides. It was always anticipated that the `load.yml` would also be relevant in this context, but it was not implemented until now.

This was a problem when an internal team needed to change the load order of two objects. Now users can do that.


Patch:
diff --git a/cumulusci/core/datasets.py b/cumulusci/core/datasets.py
index b5f21c9ca..7ebca75e1 100644
--- a/cumulusci/core/datasets.py
+++ b/cumulusci/core/datasets.py
@@ -9,6 +9,10 @@ from tempfile import TemporaryDirectory
 
 import yaml
 from simple_salesforce import Salesforce
+from snowfakery.cci_mapping_files.declaration_parser import (
+    SObjectRuleDeclaration,
+    SObjectRuleDeclarationFile,
+)
 
 from cumulusci.core import exceptions as exc
 from cumulusci.core.config import BaseProjectConfig, OrgConfig, TaskConfig
@@ -51,6 +55,11 @@ class Dataset:
     def extract_file(self) -> Path:
         return self.path / f"{self.name}.extract.yml"
 
+    @property
+    def load_rules_file(self) -> Path:
+        # to be merged into the mapping file when it is created on extract
+        return self.path / f"{self.name}.load.yml"
+
     @property
     def mapping_file(self) -> Path:
         return self.path / f"{self.name}.mapping.yml"
@@ -107,12 +116,14 @@ class Dataset:
         self,
         decls: T.Sequence[ExtractDeclaration],
         opt_in_only: T.Sequence[str] = (),
+        loading_rules: T.Sequence[SObjectRuleDeclaration] = (),
     ) -> None:
         assert isinstance(self.schema, Schema)
         mapping_data = create_load_mapping_file_from_extract_declarations(
-            decls, self.schema, opt_in_only
+            decls, self.schema, opt_in_only, loading_rules
         )
         with self.mapping_file.open("w") as f:
+            f.write(EDIT_MAPPING_WARNING)
             yaml.safe_dump(mapping_data, f, sort_keys=False)
 
     @contextmanager
@@ -150,6 +161,7 @@ class Dataset:
         logger: T.Optional[Logger] = None,
         extraction_definition: T.Optional[Path] = None,
         opt_in_only: T.Sequence[str] = (),
+        loading_rules_file: T.Optional[Path] = None,
     ):
         options = options or {}
         logger = logger or DEFAULT_LOGGER
@@ -167,9 +179,31 @@ class Dataset:
                 mapping=str(extract_mapping),
             )
             task()
-        self._save_load_mapping(list(decls.values()), opt_in_only)
+        loading_rules = self._parse_loading_rules_file(loading_rules_file)
+
+        self._save_load_mapping(list(decls.values()), opt_in_only, loading_rules)
         return task.return_values
 
+    def _parse_loading_rules_file(
+        self, loading_rules_file: T.Optional[Path]
+    ) -> T.List[SObjectRuleDeclaration]:
+        """Parse a loading rules file if provided, else try to parse
+        <datasets/dataset/dataset.load.yml>"""
+        if loading_rules_file:
+            assert loading_rules_file.exists()
+        else:
+            if self.load_rules_file.exists():
+                loading_rules_file = self.load_rules_file
+
+        if loading_rules_file:
+            with loading_rules_file.open() as f:
+                parse_result = SObjectRuleDeclarationFile.parse_from_yaml(f)
+                loading_rules = parse_result.sobject_declarations
+        else:
+            loading_rules = []
+
+        return loading_rules
+
     def _snowfakery_dataload(self, options: T.Dict, logger: Logger) -> T.Dict:
         subtask_config = TaskConfig(
             {"options": {**options, "recipe": str(self.snowfakery_recipe)}}
@@ -194,7 +228,8 @@ class Dataset:
             self._snowfakery_dataload(options, logger)
         else:  # pragma: no cover
             raise exc.BulkDataException(
-                f"Dataset has no SQL ({self.data_file}) or recipe ({self.snowfakery_recipe})"
+                f"Dataset has no SQL ({self.data_file}) "
+                "or recipe ({self.snowfakery_recipe})"
             )
 
     def _sql_dataload(self, options: T.Dict):
@@ -257,3 +292,12 @@ extract:
     OBJECTS(ALL):
         fields: FIELDS(ALL)
 """
+
+EDIT_MAPPING_WARNING = """# Editing this file is usually not recommended because it will
+# be overwritten the next time you re-capture this data.
+#
+# You can change this file's contents permanently by creating a
+# .load.yml file and re-capturing:
+#
+#  https://cumulusci.readthedocs.io/en/stable/data.html#extracting-and-loading-sample-datasets
+"""
diff --git a/cumulusci/salesforce_api/filterable_objects.py b/cumulusci/salesforce_api/filterable_objects.py
index 810307549..c93a15bd0 100644
--- a/cumulusci/salesforce_api/filterable_objects.py
+++ b/cumulusci/salesforce_api/filterable_objects.py
@@ -1,3 +1,17 @@
+# discovered by trial and error. Usually extended with a list of
+# tooling objects. i.e. any object which is both a tooling object
+# and also a "regular" sobject will be skipped.
+OPT_IN_ONLY = [
+    "ApexClass",
+    "ApexTrigger",
+    "FeedItem",
+    "Translation",
+    "WebLinkLocalization",
+    "RecordTypeLocalization",
+    "RecordType",
+    "BrandTemplate",
+]
+
 NOT_COUNTABLE = (
     "ContentDocumentLink",  # ContentDocumentLink requires a filter by a single Id on ContentDocumentId or LinkedEntityId
     "ContentFolderItem",  # Implementation restriction: ContentFolderItem requires a filter by Id or ParentContentFolderId
diff --git a/cumulusci/tasks/bulkdata/generate_mapping_utils/generate_mapping_from_declarations.py b/cumulusci/tasks/bulkdata/generate_mapping_utils/generate_mapping_from_declarations.py
index 8e7d74136..8912bb6e5 100644
--- a/cumulusci/tasks/bulkdata/generate_mapping_utils/generate_mapping_from_declarations.py
+++ b/cumulusci/tasks/bulkdata/generate_mapping_utils/generate_mapping_from_declarations.py
@@ -1,6 +1,8 @@
 import typing as T
 from itertools import chain
 
+from snowfakery.cci_mapping_files.declaration_parser import SObjectRuleDeclaration
+
 from cumulusci.salesforce_api.org_schema import Schema
 from cumulusci.tasks.bulkdata.mapping_parser import MappingStep
 from cumulusci.utils.collections import OrderedSet
@@ -23,6 +25,7 @@ def create_load_mapping_file_from_extract_declarations(
     decls: T.Sequence[ExtractDeclaration],
     schema: Schema,
     opt_in_only: T.Sequence[str] = (),
+    loading_rules: T.Sequence[SObjectRuleDeclaration] = (),
 ) -> T.Dict[str, dict]:
     """Create a mapping file from Extract declarations"""
     simplified_decls = flatten_declarations(decls, schema, opt_in_only)  # FIXME
@@ -39,7 +42,9 @@ def create_load_mapping_file_from_extract_declarations(
 
     mapping_steps = [_mapping_step(decl) for decl in simplified_decls_w_lookups]
 
-    mappings = generate_load_mapping_file(mapping_steps, intertable_dependencies, [])
+    mappings = generate_load_mapping_file(
+        mapping_steps, intertable_dependencies, loading_rules
+    )
     return mappings
 
 
diff --git a/cumulusci/tasks/bulkdata/generate_mapping_utils/load_mapping_file_generator.py b/cumulusci/tasks/bulkdata/generate_mapping_utils/load_mapping_file_generator.py
index f2a172e65..01575bf36 100644
--- a/cumulusci/tasks/bulkdata/generate_mapping_utils/load_mapping_file_generator.py
+++ b/cumulusci/tasks/bulkdata/generate_mapping_utils/load_mapping_file_generator.py
@@ -29,13 +29,13 @@ from .mapping_generator_post_processes import add_after_statements
 def generate_load_mapping_file(
     mapping_steps: T.Sequence[MappingStep],
     intertable_dependencies: OrderedSetType[SObjDependency],
-    load_declarations: T.List[SObjectRuleDeclaration] = None,
+    load_declarations: T.Sequence[SObjectRuleDeclaration] = (),
 ) -> T.Dict[str, dict]:
     """Generate a mapping file in optimal order with forward references etc.
 
     Input is a set of a tables, dependencies between them and load declarations"""
 
-    load_declarations = load_declarations or []
+    load_declarations = list(load_declarations) or []
     declared_dependencies = collect_user_specified_dependencies(load_declarations)
     table_names = OrderedSet(mapping.sf_object for mapping in mapping_steps)
     depmap = DependencyMap(
diff --git a/cumulusci/tasks/bulkdata/mapping_parser.py b/cumulusci/tasks/bulkdata/mapping_parser.py
index a0febd7fe..2b4db8ef9 100644
--- a/cumulusci/tasks/bulkdata/mapping_parser.py
+++ b/cumulusci/tasks/bulkdata/mapping_parser.py
@@ -548,9 +548,16 @@ class MappingStep(CCIDictModel):
             by_alias=by_alias, exclude_defaults=exclude_defaults, **kwargs
         )
         if fields := out.get("fields"):
+            # Convert dicts of {"Name": "Name", "Role": "Role"}
+            # (an old-fashioned syntax)
+            # into a more modern ["Name", "Role"] -type format.
             keys = list(fields.keys())
             if keys == list(fields.values()):
                 out["fields"] = keys
+
+        # flatten enum to string
+        if isinstance(out.get("api"), DataApi):
+            out["api"] = out["api"].value
         return out
 
 
diff --git a/cumulusci/tasks/sample_data/capture_sample_data.py b/cumulusci/tasks/sample_data/capture_sample_data.py
index 71ad13f39..33785263b 100644
--- a/cumulusci/tasks/sample_data/capture_sample_data.py
+++ b/cumulusci/tasks/sample_data/capture_sample_data.py
@@ -3,6 +3,7 @@ from pathlib import Path
 from cumulusci.core.config.org_config import OrgConfig
 from cumulusci.core.datasets import Dataset
 from cumulusci.core.exceptions import TaskOptionsError
+from cumulusci.salesforce_api.filterable_objects import OPT_IN_ONLY
 from cumulusci.salesforce_api.org_schema import Filters, get_org_schema
 from cumulusci.tasks.salesforce.BaseSalesforceApiTask import BaseSalesforceApiTask
 
@@ -19,7 +20,15 @@ class CaptureSampleData(BaseSalesforceApiTask):
             )
         },
         "extraction_definition": {
-            "description": "A file describing what to be extracted."
+            "description": "A file describing what to be extracted. "
+            "Defaults to `datasets/{datasetname}/{datasetname}.extract.yml` if it exists."
+        },
+        "loading_rules": {
+            "description": (
+                "Path to .load.yml file containing rules to use when loading the mapping. "
+                "Defaults to`datasets/{datasetname}/{datasetname}.load.yml` if it exists. "
+                "Multiple files can be comma separated."
+            )
         },
     }
 
@@ -35,6 +44,12 @@ class CaptureSampleData(BaseSalesforceApiTask):
             extraction_definition = Path(extraction_definition)
             if not extraction_definition.exists():
                 raise TaskOptionsError(f"Cannot find {extraction_definition}")
+
+        if loading_rules := self.options.get("loading_rules"):
+            loading_rules = Path(loading_rules)
+            if not loading_rules.exists():
+                raise TaskOptionsError(f"Cannot find {loading_rules}")
+
         with get_org_schema(
             self.sf,
             self.org_config,
@@ -53,16 +68,10 @@ class CaptureSampleData(BaseSalesforceApiTask):
             else:
                 verb = "Updated"
             opt_in_only = [f["name"] for f in self.tooling.describe()["sobjects"]]  # type: ignore
-            opt_in_only += [
-                "FeedItem",
-                "Translation",
-                "WebLinkLocalization",
-                "RecordTypeLocalization",
-                "RecordType",
-                "BrandTemplate",
-            ]
+            opt_in_only += OPT_IN_ONLY
+
             self.return_values = dataset.extract(
-                {}, self.logger, extraction_definition, opt_in_only
+                {}, self.logger, extraction_definition, opt_in_only, loading_rules
             )
             self.logger.info(f"{verb} dataset '{name}' in 'datasets/{name}'")
             return self.return_values
diff --git a/docs/data.md b/docs/data.md
index ed2c6baaf..063e3f33f 100644
--- a/docs/data.md
+++ b/docs/data.md
@@ -68,6 +68,19 @@ in a highly automated situation. It is designed to be
 used interactively, and you can control its behavior
 with an [Extract Declaration](data/extract_declarations.md) file.
 
+The extract process generates a mapping YAML file which will be used for subsequent
+loads. It has the name `datasets/<datasetname>/<datasetname>.mapping.yml`.
+It is possible to edit this file, but this may not be the best choice.
+Changes to the file can be overwritten when you capture data a second time.
+Rather than editing the file, it is preferable to create a Loading Rules
+file and then re-create the mapping file by capturing sample data again.
+
+A Loading Rules file is a file named
+`datasets/<datasetname>/<datasetname>.load.yml` which can specify instructions
+like which API to use and in which order to load objects. This file is in
+the [Loading Rules](data/loading_rules.md) format. If you create such a file
+and then re-capture sample data, the mapping file will be updated to match.
+
 ### Multiple Sample Datasets
 
 If you want different datasets for different scratch org types
@@ -919,7 +932,8 @@ Would be equivalent to `--run-until-records-loaded 700:Account` because one need
 ### Controlling the Loading Process
 
 CumulusCI's data loader has many knobs and switches that you might want
-to adjust during your load. It supports a ".load.yml" file format
+to adjust during your load. It supports a
+[".load.yml"](data/loading_rules.md) file format
 which allows you to manipulate these load settings. The simplest way to
 use this file format is to make a file in the same directory as your
 recipe with a filename that is derived from the recipe's by replacing
@@ -927,45 +941,6 @@ everything after the first "." with ".load.yml". For example, if
 your recipe is called "babka.recipe.yml" then your load file would be
 "babka.load.yml".
 
-Inside of that file you put a list of declarations in the following
-format:
-
-```yaml
-- sf_object: Account
-  api: bulk
-  bulk_mode: parallel
-```
-
-Which would specifically load accounts using the bulk API's parallel
-mode.
-
-The specific keys that you can associate with an object are:
-
--   api: "smart", "rest" or "bulk"
--   batch_size: a number
--   bulk_mode: "serial" or "parallel"
--   load_after: the name of another sobject to wait for before loading
-
-"api", "batch_size" and "bulk_mode" have the same meanings that
-they do in mapping.yml as described in [API Selection](api-selection).
-
-For example, one could force Accounts and Opportunities to load after
-Contacts:
-
-```yaml
-- sf_object: Account
-  load_after: Contact
-
-- sf_object: Opportunity
-  load_after: Contact
-```
-
-If you wish to share a loading file between multiple recipes, you can
-refer to it with the `--loading_rules` option. That will override the
-default filename (`<recipename>.load.yml`). If you want both, or any
-combination of multiple files, you can do that by listing them with
-commas between the filenames.
-
 ### Batch Sizes
 
 You can also control batch sizes with the `-o batch_size BATCHSIZE`
diff --git a/docs/data/loading_rules.md b/docs/data/loading_rules.md
new file mode 100644
index 000000000..23f2a3b95
--- /dev/null
+++ b/docs/data/loading_rules.md
@@ -0,0 +1,50 @@
+# Loading Rules
+
+CumulusCI's data loader has many knobs and switches that you might want
+to adjust during your load. It supports a ".load.yml" file format
+which allows you to manipulate these load settings. The simplest way to
+use this file format is to make a file in the same directory as your
+recipe or dataset with a filename that is derived from the recipe's by replacing
+everything after the first "." with ".load.yml". For example, if
+your recipe is called "babka.recipe.yml" then your load file would be
+"babka.load.yml". Similarly if you have a dataset called
+`datasets/ds/ds.dataset.sql` then you could create `datasets/ds/ds.load.yml`.
+
+Inside of that file you put a list of declarations in the following
+format:
+
+```yaml
+- sf_object: Account
+  api: bulk
+  bulk_mode: parallel
+```
+
+Which would specifically load accounts using the bulk API's parallel
+mode.
+
+The specific keys that you can associate with an object are:
+
+-   api: "smart", "rest" or "bulk"
+-   batch_size: a number
+-   bulk_mode: "serial" or "parallel"
+-   load_after: the name of another sobject to wait for before loading
+
+"api", "batch_size" and "bulk_mode" have the same meanings that
+they do in mapping.yml as described in [API Selection](api-selection).
+
+For example, one could force Accounts and Opportunities to load after
+Contacts:
+
+```yaml
+- sf_object: Account
+  load_after: Contact
+
+- sf_object: Opportunity
+  load_after: Contact
+```
+
+If you wish to share a loading file between multiple recipes or
+datasets, you can refer to it with the `--loading_rules` option.
+That will override the default filename (`<recipename>.load.yml`).
+If you want both, or any combination of multiple files, you can do
+that by listing them with commas between the filenames.


Test patch:
diff --git a/cumulusci/core/tests/test_datasets_e2e.py b/cumulusci/core/tests/test_datasets_e2e.py
index 3033e5c06..5c9bdfb86 100644
--- a/cumulusci/core/tests/test_datasets_e2e.py
+++ b/cumulusci/core/tests/test_datasets_e2e.py
@@ -3,6 +3,7 @@ from contextlib import contextmanager
 from pathlib import Path
 from shutil import rmtree
 from tempfile import TemporaryDirectory
+from typing import Any
 from unittest.mock import patch
 
 import pytest
@@ -113,6 +114,9 @@ class TestDatasetsE2E:
         timer.checkpoint("Before create")
         # Create a dataset directory and associated files
         dataset.create()
+        load_rules = [{"sf_object": "Account", "api": "rest"}]
+        dataset.load_rules_file.write_text(yaml.safe_dump(load_rules))
+
         timer.checkpoint("After create")
 
         # Read selected schema objects and fields into a T.Dict[str, T.List[str]]
@@ -129,7 +133,8 @@ class TestDatasetsE2E:
         dataset.update_schema_subset(objs)
         timer.checkpoint("Updated Subset")
 
-        # Read full schema in a Datastructure like {"Account": {"Name": True, "Description": False}}
+        # Read full schema in a Datastructure like
+        # {"Account": {"Name": True, "Description": False}}
         objs = dataset.read_which_fields_selected()
         assert objs["Account"]["Name"] is True
         assert objs["Account"]["Description"] is False
@@ -243,23 +248,57 @@ class TestDatasetsE2E:
 
             dataset.create()
             with TemporaryDirectory() as t:
-                def_file = Path(t) / "test_extract_definition"
-                with open(def_file, "w") as f:
-                    yaml.safe_dump({"extract": {"Account": {"fields": "Name"}}}, f)
+
+                def write_yaml(filename: str, json: Any):
+                    full_path = Path(t) / filename
+                    with open(full_path, "w") as f:
+                        yaml.safe_dump(json, f)
+                    return full_path
+
+                extract_def = write_yaml(
+                    "test_extract_definition",
+                    {
+                        "extract": {
+                            "Account": {"fields": ["Name"]},
+                            "Contact": {
+                                "fields": ["FirstName", "LastName", "AccountId"]
+                            },
+                        }
+                    },
+                )
+                loading_rules = write_yaml(
+                    "loading_rules.load.yml",
+                    [{"sf_object": "Account", "load_after": "Contact"}],
+                )
 
                 # Don't actually extract data.
                 with patch("cumulusci.tasks.bulkdata.ExtractData._run_task"):
-                    dataset.extract(extraction_definition=def_file)
+                    dataset.extract(
+                        extraction_definition=extract_def,
+                        loading_rules_file=loading_rules,
+                    )
             with open(dataset.mapping_file) as p:
                 actual = yaml.safe_load(p)
                 expected = {
+                    "Insert Contact": {
+                        "sf_object": "Contact",
+                        "table": "Contact",
+                        "fields": ["FirstName", "LastName"],
+                        "lookups": {
+                            "AccountId": {
+                                "table": "Account",
+                                "key_field": "AccountId",
+                                "after": "Insert Account",
+                            }
+                        },
+                    },
                     "Insert Account": {
                         "sf_object": "Account",
                         "table": "Account",
                         "fields": ["Name"],
-                    }
+                    },
                 }
-                assert actual == expected, actual
+                assert tuple(actual.items()) == tuple(expected.items()), actual.items()
 
             if dataset.path.exists():
                 rmtree(dataset.path)
@@ -289,7 +328,6 @@ class TestLoadDatasets:
         assert called
 
     def test_dataset_with_no_data_or_recipe(self, sf, project_config, org_config):
-
         with setup_test(org_config), Dataset(
             "fxoyoxz", project_config, sf, org_config, schema=None
         ) as dataset, pytest.raises(BulkDataException, match="fxoyoxz"):
diff --git a/cumulusci/tasks/bulkdata/generate_mapping_utils/tests/test_generate_load_mapping_from_declarations.py b/cumulusci/tasks/bulkdata/generate_mapping_utils/tests/test_generate_load_mapping_from_declarations.py
index 5996d819f..9b163bfe3 100644
--- a/cumulusci/tasks/bulkdata/generate_mapping_utils/tests/test_generate_load_mapping_from_declarations.py
+++ b/cumulusci/tasks/bulkdata/generate_mapping_utils/tests/test_generate_load_mapping_from_declarations.py
@@ -44,6 +44,45 @@ class TestGenerateLoadMappingFromDeclarations:
                 }
             }
 
+    def test_generate_mapping_from_both_kinds_of_declarations(self, org_config):
+        declarations = [
+            ExtractDeclaration(sf_object="Account", fields=["Name", "Description"]),
+            ExtractDeclaration(sf_object="Contact", fields=["FirstName", "LastName"]),
+        ]
+        object_counts = {"Account": 10, "Contact": 1, "Case": 5}
+        obj_describes = (
+            describe_for("Account"),
+            describe_for("Contact"),
+            describe_for("Case"),
+        )
+        with _fake_get_org_schema(
+            org_config,
+            obj_describes,
+            object_counts,
+            filters=[],
+            include_counts=True,
+        ) as schema:
+            loading_rules = [
+                SObjectRuleDeclaration(sf_object="Account", load_after="Contact")
+            ]
+            mf = create_load_mapping_file_from_extract_declarations(
+                declarations, schema, (), loading_rules
+            )
+            assert tuple(mf.items()) == tuple(
+                {
+                    "Insert Contact": {
+                        "sf_object": "Contact",
+                        "table": "Contact",
+                        "fields": ["FirstName", "LastName"],
+                    },
+                    "Insert Account": {
+                        "sf_object": "Account",
+                        "table": "Account",
+                        "fields": ["Name", "Description"],
+                    },
+                }.items()
+            )
+
     def test_generate_load_mapping_from_declarations__lookups(self, org_config):
         declarations = [
             ExtractDeclaration(sf_object="Account", fields=["Name", "Description"]),
@@ -162,7 +201,7 @@ class TestGenerateLoadMappingFromDeclarations:
             },
         }, mf
 
-    def test_generate_load_mapping__with_userts(self, org_config):
+    def test_generate_load_mapping__with_upserts(self, org_config):
 
         mapping_steps = [
             MappingStep(sf_object="Account"),
@@ -184,7 +223,7 @@ class TestGenerateLoadMappingFromDeclarations:
         mf = generate_load_mapping_file(
             mapping_steps,
             intertable_dependencies=set(),
-            load_declarations=None,
+            load_declarations=(),
         )
         assert mf == {
             "Insert Account": {
diff --git a/cumulusci/tasks/sample_data/test_capture_sample_data.py b/cumulusci/tasks/sample_data/test_capture_sample_data.py
index 8db777374..d4819de50 100644
--- a/cumulusci/tasks/sample_data/test_capture_sample_data.py
+++ b/cumulusci/tasks/sample_data/test_capture_sample_data.py
@@ -60,7 +60,7 @@ class TestCaptureDatasets:
             Dataset.assert_any_call("default", mock.ANY, mock.ANY, org_config, mock.ANY)
             # and extracted
             Dataset().__enter__().extract.assert_called_with(
-                {}, task.logger, None, mock.ANY
+                {}, task.logger, None, mock.ANY, None
             )
 
     @mock.patch("cumulusci.tasks.sample_data.capture_sample_data.Dataset")
@@ -75,12 +75,17 @@ class TestCaptureDatasets:
             with open(extraction_definition, "w") as f:
                 f.write("")  # Doesn't matter. We won't parse it
 
+            loading_rules = Path(t) / "test_load_definition"
+            with open(loading_rules, "w") as f:
+                f.write("[]")  # Doesn't matter. We won't parse it
+
             Dataset().__enter__().path.exists.return_value = False
             task = create_task(
                 CaptureSampleData,
                 {
                     "dataset": "mydataset",
                     "extraction_definition": extraction_definition,
+                    "loading_rules": loading_rules,
                 },
             )
             task()
@@ -91,7 +96,7 @@ class TestCaptureDatasets:
             Dataset().__enter__().create.assert_called_with()
             # and extracted
             Dataset().__enter__().extract.assert_called_with(
-                {}, task.logger, extraction_definition, mock.ANY
+                {}, task.logger, extraction_definition, mock.ANY, loading_rules
             )
 
     @mock.patch("cumulusci.tasks.sample_data.capture_sample_data.Dataset")