From 28926b0f80ae5fb4228d5ba2743ec146f69c7963 Mon Sep 17 00:00:00 2001
From: tjchadaga <85581939+tjchadaga@users.noreply.github.com>
Date: Fri, 8 Jul 2022 16:54:38 -0700
Subject: [PATCH] Add 'traffic_shift_away' option to config load_minigraph
 (#2240)

---
 config/main.py           | 10 +++++++++-
 doc/Command-Reference.md |  4 +++-
 tests/config_test.py     | 28 ++++++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/config/main.py b/config/main.py
index 8a4c595562..f95c64c2b8 100644
--- a/config/main.py
+++ b/config/main.py
@@ -1685,8 +1685,9 @@ def load_mgmt_config(filename):
 @click.option('-y', '--yes', is_flag=True, callback=_abort_if_false,
                 expose_value=False, prompt='Reload config from minigraph?')
 @click.option('-n', '--no_service_restart', default=False, is_flag=True, help='Do not restart docker services')
+@click.option('-t', '--traffic_shift_away', default=False, is_flag=True, help='Keep device in maintenance with TSA')
 @clicommon.pass_db
-def load_minigraph(db, no_service_restart):
+def load_minigraph(db, no_service_restart, traffic_shift_away):
     """Reconfigure based on minigraph."""
     log.log_info("'load_minigraph' executing...")
 
@@ -1756,6 +1757,13 @@ def load_minigraph(db, no_service_restart):
                 cfggen_namespace_option = " -n {}".format(namespace)
             clicommon.run_command(db_migrator + ' -o set_version' + cfggen_namespace_option)
 
+    # Keep device isolated with TSA 
+    if traffic_shift_away:
+        clicommon.run_command("TSA", display_cmd=True)
+        if os.path.isfile(DEFAULT_GOLDEN_CONFIG_DB_FILE):
+            log.log_warning("Golden configuration may override System Maintenance state. Please execute TSC to check the current System mode")
+            click.secho("[WARNING] Golden configuration may override Traffic-shift-away state. Please execute TSC to check the current System mode")
+
     # Load golden_config_db.json
     if os.path.isfile(DEFAULT_GOLDEN_CONFIG_DB_FILE):
         override_config_by(DEFAULT_GOLDEN_CONFIG_DB_FILE)
diff --git a/doc/Command-Reference.md b/doc/Command-Reference.md
index 2c873235e1..012834b017 100644
--- a/doc/Command-Reference.md
+++ b/doc/Command-Reference.md
@@ -5149,9 +5149,11 @@ When user specifies the optional argument "-n" or "--no-service-restart", this c
 running on the device. One use case for this option is during boot time when config-setup service loads minigraph configuration and there is no services
 running on the device.
 
+When user specifies the optional argument "-t" or "--traffic-shift-away", this command executes TSA command at the end to ensure the device remains in maintenance after loading minigraph.
+
 - Usage:
   ```
-  config load_minigraph [-y|--yes] [-n|--no-service-restart]
+  config load_minigraph [-y|--yes] [-n|--no-service-restart] [-t|--traffic-shift-away]
   ```
 
 - Example:
diff --git a/tests/config_test.py b/tests/config_test.py
index ca06900817..18fa251e9b 100644
--- a/tests/config_test.py
+++ b/tests/config_test.py
@@ -435,6 +435,34 @@ def is_file_side_effect(filename):
             assert result.exit_code == 0
             assert expected_output in result.output
 
+    def test_load_minigraph_with_traffic_shift_away(self, get_cmd_module):
+        with mock.patch("utilities_common.cli.run_command", mock.MagicMock(side_effect=mock_run_command_side_effect)) as mock_run_command:
+            (config, show) = get_cmd_module
+            runner = CliRunner()
+            result = runner.invoke(config.config.commands["load_minigraph"], ["-ty"])
+            print(result.exit_code)
+            print(result.output)
+            traceback.print_tb(result.exc_info[2])
+            assert result.exit_code == 0
+            assert "TSA" in result.output
+
+    def test_load_minigraph_with_traffic_shift_away_with_golden_config(self, get_cmd_module):
+        with mock.patch("utilities_common.cli.run_command", mock.MagicMock(side_effect=mock_run_command_side_effect)) as mock_run_command:
+            def is_file_side_effect(filename):
+                return True if 'golden_config' in filename else False
+            with mock.patch('os.path.isfile', mock.MagicMock(side_effect=is_file_side_effect)):
+                (config, show) = get_cmd_module
+                db = Db()
+                golden_config = {}
+                runner = CliRunner()
+                result = runner.invoke(config.config.commands["load_minigraph"], ["-ty"])
+                print(result.exit_code)
+                print(result.output)
+                traceback.print_tb(result.exc_info[2])
+                assert result.exit_code == 0
+                assert "TSA" in result.output
+                assert "[WARNING] Golden configuration may override Traffic-shift-away state" in result.output
+
     @classmethod
     def teardown_class(cls):
         os.environ['UTILITIES_UNIT_TESTING'] = "0"