Add 'traffic_shift_away' option to config load_minigraph

config/main.py

@@ -1641,8 +1641,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...")
 
@@ -1712,6 +1713,10 @@ def load_minigraph(db, no_service_restart):
     if os.path.isfile(DEFAULT_GOLDEN_CONFIG_DB_FILE):
         override_config_by(DEFAULT_GOLDEN_CONFIG_DB_FILE)
 
+    # Keep device in maintenance with TSA 
+    if traffic_shift_away:
+        clicommon.run_command("TSA", display_cmd=True)
+
     # We first run "systemctl reset-failed" to remove the "failed"
     # status from all services before we attempt to restart them
     if not no_service_restart:

doc/Command-Reference.md

@@ -5106,9 +5106,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:

tests/config_test.py

@@ -356,6 +356,17 @@ 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
+
     @classmethod
     def teardown_class(cls):
         os.environ['UTILITIES_UNIT_TESTING'] = "0"