From 7b71db1a70e60cd72492bf0c7499dfc4740b2754 Mon Sep 17 00:00:00 2001
From: Paul Prescod <pprescod@salesforce.com>
Date: Thu, 28 Oct 2021 09:25:52 -0700
Subject: [PATCH] Parse user_channel declarations (#520)

---
 snowfakery/api.py                             |  3 +-
 .../cci_mapping_files/declaration_parser.py   | 52 ++++++++++++++++---
 tests/cci/test_declaration_parser.py          | 17 +++++-
 tests/cci/user_channels.load.yml              | 17 ++++++
 4 files changed, 80 insertions(+), 9 deletions(-)
 create mode 100644 tests/cci/user_channels.load.yml

diff --git a/snowfakery/api.py b/snowfakery/api.py
index 5441ad83..7524a0b2 100644
--- a/snowfakery/api.py
+++ b/snowfakery/api.py
@@ -302,7 +302,8 @@ def gather_declarations(yaml_file, load_declarations):
         declarations = []
         for declfile in load_declarations:
             with open_file_like(declfile, "r") as (path, f):
-                declarations.extend(SObjectRuleDeclarationFile.parse_from_yaml(f))
+                decls = SObjectRuleDeclarationFile.parse_from_yaml(f)
+                declarations.extend(decls.sobject_declarations)
 
         unified_declarations = unify(declarations)
     else:
diff --git a/snowfakery/cci_mapping_files/declaration_parser.py b/snowfakery/cci_mapping_files/declaration_parser.py
index 86f6dcb0..83ec262c 100644
--- a/snowfakery/cci_mapping_files/declaration_parser.py
+++ b/snowfakery/cci_mapping_files/declaration_parser.py
@@ -1,13 +1,11 @@
-from pathlib import Path
 import typing as T
-
-from datetime import date
 from collections import defaultdict
+from datetime import date
+from pathlib import Path
 
 import yaml
-
+from pydantic import BaseModel, Extra, validator
 from typing_extensions import Literal
-from pydantic import BaseModel, validator, Extra
 
 
 class AtomicDecl(T.NamedTuple):
@@ -105,8 +103,31 @@ def as_mapping(self):
 }
 
 
+class ChannelDeclaration(BaseModel):
+    "Channel declarations are only of relevance to Salesforce employees"
+    user: str
+    recipe_options: T.Dict[str, T.Any] = None
+    num_generators: int = None
+    num_loaders: int = None
+
+    class Config:
+        extra = Extra.forbid
+
+
+class ChannelDeclarationList(BaseModel):
+    "Channel declarations are only of relevance to Salesforce employees"
+    user_channels: T.List[ChannelDeclaration]
+
+
+class LoadDeclarationsTuple(T.NamedTuple):
+    sobject_declarations: T.List[SObjectRuleDeclaration]
+    channel_declarations: T.List[
+        ChannelDeclaration
+    ]  # Channel declarations are only of relevance to Salesforce employees
+
+
 class SObjectRuleDeclarationFile(BaseModel):
-    __root__: T.List[SObjectRuleDeclaration]
+    __root__: T.List[T.Union[ChannelDeclarationList, SObjectRuleDeclaration]]
 
     @classmethod
     def parse_from_yaml(cls, f: T.Union[Path, T.TextIO]):
@@ -117,7 +138,24 @@ def parse_from_yaml(cls, f: T.Union[Path, T.TextIO]):
         else:
             data = yaml.safe_load(f)
 
-        return cls.parse_obj(data).__root__
+        sobject_decls = [
+            obj
+            for obj in cls.parse_obj(data).__root__
+            if isinstance(obj, SObjectRuleDeclaration)
+        ]
+        channel_decls = [
+            obj
+            for obj in cls.parse_obj(data).__root__
+            if isinstance(obj, ChannelDeclarationList)
+        ]
+        if len(channel_decls) > 1:
+            raise AssertionError("Only one channel declaration list allowed per file.")
+        elif len(channel_decls) == 1:
+            channels = channel_decls[0].user_channels
+        else:
+            channels = []
+
+        return LoadDeclarationsTuple(sobject_decls, channels)
 
 
 def atomize_decls(decls: T.Sequence[SObjectRuleDeclaration]):
diff --git a/tests/cci/test_declaration_parser.py b/tests/cci/test_declaration_parser.py
index 0377d240..f8a1e89c 100644
--- a/tests/cci/test_declaration_parser.py
+++ b/tests/cci/test_declaration_parser.py
@@ -164,4 +164,19 @@ def test_parse_from_open_file(self):
 
     def test_parse_from_path(self):
         sample_yaml = Path(__file__).parent / "mapping_mixins-override.load.yml"
-        SObjectRuleDeclarationFile.parse_from_yaml(sample_yaml)
+        assert SObjectRuleDeclarationFile.parse_from_yaml(sample_yaml)
+
+
+#  Channel declarations are only of relevance to Salesforce employees
+class TestUserChannels:
+    def test_user_channel_declarations(self):
+        sample_yaml = Path(__file__).parent / "user_channels.load.yml"
+        channels = SObjectRuleDeclarationFile.parse_from_yaml(
+            sample_yaml
+        ).channel_declarations
+        assert len(channels) == 4
+        assert channels[0].user == "admin_1_org_alias"
+        assert channels[1].recipe_options == {"a": "b"}
+        assert channels[2].num_generators == 5
+        assert channels[3].user == "user_3_org_alias"
+        assert channels[3].num_loaders == 11
diff --git a/tests/cci/user_channels.load.yml b/tests/cci/user_channels.load.yml
new file mode 100644
index 00000000..13b9dba7
--- /dev/null
+++ b/tests/cci/user_channels.load.yml
@@ -0,0 +1,17 @@
+# Channel declarations are only of relevance to Salesforce employees
+- sf_object: Account
+  api: REST
+  priority: high
+
+- sf_object: Opportunity
+  api: bulk
+
+- user_channels:
+    - user: admin_1_org_alias
+    - user: user_1_org_alias
+      recipe_options:
+        a: b
+    - user: user_2_org_alias
+      num_generators: 5
+    - user: user_3_org_alias
+      num_loaders: 11