[auto_ts] Enable register/de-register auto_ts config for APP Extension

sonic_package_manager/service_creator/feature.py

@@ -1,9 +1,10 @@
 #!/usr/bin/env python
 
 """ This module implements new feature registration/de-registration in SONiC system. """
-
+import copy
 from typing import Dict, Type
 
+from sonic_package_manager.logger import log
 from sonic_package_manager.manifest import Manifest
 from sonic_package_manager.service_creator.sonic_db import SonicDB
 
@@ -15,6 +16,14 @@
     'set_owner': 'local'
 }
 
+AUTO_TS_GLOBAL = "AUTO_TECHSUPPORT"
+AUTO_TS_FEATURE = "AUTO_TECHSUPPORT_FEATURE"
+CFG_STATE = "state"
+# TODO: Enable available_mem_threshold once the mem_leak_auto_ts feature is available
+DEFAULT_AUTO_TS_FEATURE_CONFIG = {
+    'state': 'disabled',
+    'rate_limit_interval': '600'
+}
 
 def is_enabled(cfg):
     return cfg.get('state', 'disabled').lower() == 'enabled'
@@ -25,8 +34,11 @@ def is_multi_instance(cfg):
 
 
 class FeatureRegistry:
-    """ FeatureRegistry class provides an interface to
-    register/de-register new feature persistently. """
+    """ 1) FeatureRegistry class provides an interface to
+    register/de-register new feature tables persistently.
+        2) Writes persistent configuration to FEATURE & 
+    AUTO_TECHSUPPORT_FEATURE tables
+    """
 
     def __init__(self, sonic_db: Type[SonicDB]):
         self._sonic_db = sonic_db
@@ -60,6 +72,9 @@ def register(self,
             new_cfg = {**new_cfg, **non_cfg_entries}
 
             conn.set_entry(FEATURE, name, new_cfg)
+
+        if self.register_auto_ts(name):
+            log.info(f'{name} entry is added to {AUTO_TS_FEATURE} table')
 
     def deregister(self, name: str):
         """ Deregister feature by name.
@@ -73,6 +88,7 @@ def deregister(self, name: str):
         db_connetors = self._sonic_db.get_connectors()
         for conn in db_connetors:
             conn.set_entry(FEATURE, name, None)
+            conn.set_entry(AUTO_TS_FEATURE, name, None)
 
     def update(self,
                old_manifest: Manifest,
@@ -103,6 +119,9 @@ def update(self,
             new_cfg = {**new_cfg, **non_cfg_entries}
 
             conn.set_entry(FEATURE, new_name, new_cfg)
+
+        if self.register_auto_ts(new_name, old_name):
+            log.info(f'{new_name} entry is added to {AUTO_TS_FEATURE} table')
 
     def is_feature_enabled(self, name: str) -> bool:
         """ Returns whether the feature is current enabled
@@ -123,6 +142,46 @@ def get_multi_instance_features(self):
         features = conn.get_table(FEATURE)
         return [feature for feature, cfg in features.items() if is_multi_instance(cfg)]
 
+    def infer_auto_ts_capability(self, init_cfg_conn):
+        """ Determine whether to enable/disable the state for new feature
+        AUTO_TS provides a compile-time knob to enable/disable this feature
+        Default State for the new feature follows the decision made at compile time.
+
+        Args:
+            init_cfg_conn: PersistentConfigDbConnector for init_cfg.json
+        Returns:
+            Capability: Tuple: (bool, ["enabled", "disabled"])
+        """
+        cfg = init_cfg_conn.get_entry(AUTO_TS_GLOBAL, "GLOBAL")
+        default_state = cfg.get(CFG_STATE, "")
+        if not default_state:
+            return (False, "disabled")
+        else:
+            return (True, default_state)
+
+    def register_auto_ts(self, new_name, old_name=None):
+        """ Registers auto_ts feature
+        """
+        # Infer and update default config
+        init_cfg_conn = self._sonic_db.get_initial_db_connector()
+        def_cfg = DEFAULT_AUTO_TS_FEATURE_CONFIG.copy()
+        (auto_ts_add_cfg, auto_ts_state) = self.infer_auto_ts_capability(init_cfg_conn)
+        def_cfg['state'] = auto_ts_state
+
+        if not auto_ts_add_cfg:
+            log.debug("Skip adding AUTO_TECHSUPPORT_FEATURE table because no AUTO_TECHSUPPORT|GLOBAL entry is found")
+            return False
+
+        for conn in self._sonic_db.get_connectors():
+            new_cfg = copy.deepcopy(def_cfg)
+            if old_name:
+                current_cfg = conn.get_entry(AUTO_TS_FEATURE, old_name)
+                conn.set_entry(AUTO_TS_FEATURE, old_name, None)
+                new_cfg.update(current_cfg)
+
+            conn.set_entry(AUTO_TS_FEATURE, new_name, new_cfg)
+        return True
+
     @staticmethod
     def get_default_feature_entries(state=None, owner=None) -> Dict[str, str]:
         """ Get configurable feature table entries:

tests/sonic_package_manager/test_service_creator.py

@@ -2,6 +2,7 @@
 
 import os
 import copy
+import unittest
 from unittest.mock import Mock, MagicMock, call
 
 import pytest
@@ -205,6 +206,7 @@ def test_feature_registration(mock_sonic_db, manifest):
     mock_connector = Mock()
     mock_connector.get_entry = Mock(return_value={})
     mock_sonic_db.get_connectors = Mock(return_value=[mock_connector])
+    mock_sonic_db.get_initial_db_connector = Mock(return_value=mock_connector)
     feature_registry = FeatureRegistry(mock_sonic_db)
     feature_registry.register(manifest)
     mock_connector.set_entry.assert_called_with('FEATURE', 'test', {
@@ -258,6 +260,7 @@ def test_feature_registration_with_timer(mock_sonic_db, manifest):
     mock_connector = Mock()
     mock_connector.get_entry = Mock(return_value={})
     mock_sonic_db.get_connectors = Mock(return_value=[mock_connector])
+    mock_sonic_db.get_initial_db_connector = Mock(return_value=mock_connector)
     feature_registry = FeatureRegistry(mock_sonic_db)
     feature_registry.register(manifest)
     mock_connector.set_entry.assert_called_with('FEATURE', 'test', {
@@ -275,6 +278,7 @@ def test_feature_registration_with_non_default_owner(mock_sonic_db, manifest):
     mock_connector = Mock()
     mock_connector.get_entry = Mock(return_value={})
     mock_sonic_db.get_connectors = Mock(return_value=[mock_connector])
+    mock_sonic_db.get_initial_db_connector = Mock(return_value=mock_connector)
     feature_registry = FeatureRegistry(mock_sonic_db)
     feature_registry.register(manifest, owner='kube')
     mock_connector.set_entry.assert_called_with('FEATURE', 'test', {
@@ -286,3 +290,107 @@ def test_feature_registration_with_non_default_owner(mock_sonic_db, manifest):
         'has_global_scope': 'True',
         'has_timer': 'False',
     })
+
+
+class AutoTSHelp:
+    """ Helper class for Auto TS Feature Registry Tests
+    """
+    GLOBAL_STATE = {}
+
+    @classmethod
+    def get_entry(cls, table, key):
+        if table == "AUTO_TECHSUPPORT" and key == "GLOBAL":
+            return AutoTSHelp.GLOBAL_STATE
+        elif table == "AUTO_TECHSUPPORT_FEATURE" and key == "test":
+            return {"state" : "enabled", "rate_limit_interval" : "600"}
+        else:
+            return {}
+
+    @classmethod
+    def get_entry_running_cfg(cls, table, key):
+        if table == "AUTO_TECHSUPPORT_FEATURE" and key == "test":
+            return {"state" : "disabled", "rate_limit_interval" : "1000"}
+        else:
+            return {}
+
+
+def test_auto_ts_global_disabled(mock_sonic_db, manifest):
+    mock_init_cfg = Mock()
+    AutoTSHelp.GLOBAL_STATE = {"state" : "disabled"}
+    mock_init_cfg.get_entry = Mock(side_effect=AutoTSHelp.get_entry)
+    mock_sonic_db.get_connectors = Mock(return_value=[mock_init_cfg])
+    mock_sonic_db.get_initial_db_connector = Mock(return_value=mock_init_cfg)
+    feature_registry = FeatureRegistry(mock_sonic_db)
+    feature_registry.register(manifest)
+    mock_init_cfg.set_entry.assert_any_call("AUTO_TECHSUPPORT_FEATURE", "test", {
+            "state" : "disabled",
+            "rate_limit_interval" : "600"
+        }
+    )
+
+
+def test_auto_ts_global_enabled(mock_sonic_db, manifest):
+    mock_init_cfg = Mock()
+    AutoTSHelp.GLOBAL_STATE = {"state" : "enabled"}
+    mock_init_cfg.get_entry = Mock(side_effect=AutoTSHelp.get_entry)
+    mock_sonic_db.get_connectors = Mock(return_value=[mock_init_cfg])
+    mock_sonic_db.get_initial_db_connector = Mock(return_value=mock_init_cfg)
+    feature_registry = FeatureRegistry(mock_sonic_db)
+    feature_registry.register(manifest)
+    mock_init_cfg.set_entry.assert_any_call("AUTO_TECHSUPPORT_FEATURE", "test", {
+            "state" : "enabled",
+            "rate_limit_interval" : "600"
+        }
+    )
+
+
+def test_auto_ts_deregister(mock_sonic_db):
+    mock_connector = Mock()
+    mock_sonic_db.get_connectors = Mock(return_value=[mock_connector])
+    feature_registry = FeatureRegistry(mock_sonic_db)
+    feature_registry.deregister("test")
+    mock_connector.set_entry.assert_any_call("AUTO_TECHSUPPORT_FEATURE", "test", None)
+
+
+def test_auto_ts_feature_update_flow(mock_sonic_db, manifest):
+    new_manifest = copy.deepcopy(manifest)
+    new_manifest['service']['name'] = 'test_new'
+    new_manifest['service']['delayed'] = True
+
+    AutoTSHelp.GLOBAL_STATE = {"state" : "enabled"}
+    # Mock init_cfg connector
+    mock_init_cfg = Mock()
+    mock_init_cfg.get_entry = Mock(side_effect=AutoTSHelp.get_entry)
+
+    # Mock running/peristent cfg connector
+    mock_other_cfg = Mock()
+    mock_other_cfg.get_entry = Mock(side_effect=AutoTSHelp.get_entry_running_cfg)
+
+    # Setup sonic_db class
+    mock_sonic_db.get_connectors = Mock(return_value=[mock_init_cfg, mock_other_cfg])
+    mock_sonic_db.get_initial_db_connector = Mock(return_value=mock_init_cfg)
+
+    feature_registry = FeatureRegistry(mock_sonic_db)
+    feature_registry.update(manifest, new_manifest)
+
+    mock_init_cfg.set_entry.assert_has_calls(
+        [
+            call("AUTO_TECHSUPPORT_FEATURE", "test", None),
+            call("AUTO_TECHSUPPORT_FEATURE", "test_new", {
+                    "state" : "enabled",
+                    "rate_limit_interval" : "600"
+                })
+        ],
+        any_order = True
+    )
+
+    mock_other_cfg.set_entry.assert_has_calls(
+        [
+            call("AUTO_TECHSUPPORT_FEATURE", "test", None),
+            call("AUTO_TECHSUPPORT_FEATURE", "test_new", {
+                    "state" : "disabled",
+                    "rate_limit_interval" : "1000"
+                })
+        ],
+        any_order = True
+    )