Merge pull request #52221 from garethgreenaway/pip_state_mod_aggregate

[develop] First pass at adding mod_aggregate support to the pip_state module.

salt/states/pip_state.py

@@ -1065,3 +1065,51 @@ def uptodate(name,
         ret['comment'] = 'Upgrade failed.'
 
     return ret
+
+
+def mod_aggregate(low, chunks, running):
+    '''
+    The mod_aggregate function which looks up all packages in the available
+    low chunks and merges them into a single pkgs ref in the present low data
+    '''
+    pkgs = []
+    pkg_type = None
+    agg_enabled = [
+        'installed',
+        'removed',
+    ]
+    if low.get('fun') not in agg_enabled:
+        return low
+    for chunk in chunks:
+        tag = __utils__['state.gen_tag'](chunk)
+        if tag in running:
+            # Already ran the pkg state, skip aggregation
+            continue
+        if chunk.get('state') == 'pip':
+            if '__agg__' in chunk:
+                continue
+            # Check for the same function
+            if chunk.get('fun') != low.get('fun'):
+                continue
+            # Check first if 'sources' was passed so we don't aggregate pkgs
+            # and sources together.
+            if pkg_type is None:
+                pkg_type = 'pkgs'
+            if pkg_type == 'pkgs':
+                # Pull out the pkg names!
+                if 'pkgs' in chunk:
+                    pkgs.extend(chunk['pkgs'])
+                    chunk['__agg__'] = True
+                elif 'name' in chunk:
+                    version = chunk.pop('version', None)
+                    if version is not None:
+                        pkgs.append({chunk['name']: version})
+                    else:
+                        pkgs.append(chunk['name'])
+                    chunk['__agg__'] = True
+    if pkg_type is not None and pkgs:
+        if pkg_type in low:
+            low[pkg_type].extend(pkgs)
+        else:
+            low[pkg_type] = pkgs
+    return low

tests/unit/states/test_pip_state.py

@@ -18,6 +18,7 @@
 # Import salt libs
 import salt.states.pip_state as pip_state
 import salt.utils.versions
+from salt.utils.odict import OrderedDict
 
 # Import 3rd-party libs
 try:
@@ -304,3 +305,45 @@ def test_install_in_editable_mode(self):
                 'successfully installed',
                 {'test': ret}
             )
+
+    def test_mod_aggregate(self):
+        '''
+        Test to mod_aggregate function
+        '''
+
+        low = OrderedDict([('state', 'pip'),
+                           ('name', 'ipython'),
+                           ('__sls__', 'test.test_pip'),
+                           ('__env__', 'base'),
+                           ('__id__', 'ipython'),
+                           ('order', 10000),
+                           ('fun', 'installed')])
+        chunks = [OrderedDict([('state', 'pip'),
+                               ('name', 'ipython'),
+                               ('__sls__', 'test.test_pip'),
+                               ('__env__', 'base'),
+                               ('__id__', 'ipython'),
+                               ('order', 10000),
+                               ('fun', 'installed')]),
+                  OrderedDict([('state', 'pip'),
+                               ('name', 'pylint'),
+                               ('__sls__', 'test.test_pip'),
+                               ('__env__', 'base'),
+                               ('__id__', 'pylint'),
+                               ('order', 10001),
+                               ('fun', 'installed')])]
+        running = {}
+        expected_low = OrderedDict([('state', 'pip'),
+                                    ('name', 'ipython'),
+                                    ('__sls__', 'test.test_pip'),
+                                    ('__env__', 'base'),
+                                    ('__id__', 'ipython'),
+                                    ('order', 10000),
+                                    ('fun', 'installed'),
+                                    ('pkgs', ['ipython', 'pylint'])])
+
+        mock_tag = MagicMock(side_effect=['pip_|-ipython_|-ipython_|-installed',
+                                          'pip_|-pylint_|-pylint_|-installed'])
+        with patch.dict(pip_state.__utils__, {'state.gen_tag': mock_tag}):
+            self.assertDictEqual(pip_state.mod_aggregate(low, chunks, running),
+                                 expected_low)