From 9c1d489ca5cd849020610c4a57664c675c64d321 Mon Sep 17 00:00:00 2001
From: ganglv <88995770+ganglyu@users.noreply.github.com>
Date: Mon, 25 Dec 2023 16:05:16 +0800
Subject: [PATCH] Fix database initialization for db_migrator (#3100)

What I did
db_migrator failed to initialize SonicDBConfig, and I fix this issue.

How I did it
If SonicDBConfig is already initialized, do not invoke initialize() again.

How to verify it
Run unit test, and verified on DUT.
---
 scripts/db_migrator.py    |  8 +++-----
 tests/db_migrator_test.py | 18 +++++++++++++++++-
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/scripts/db_migrator.py b/scripts/db_migrator.py
index f3ffe0197d..a48da00f5b 100755
--- a/scripts/db_migrator.py
+++ b/scripts/db_migrator.py
@@ -8,9 +8,10 @@
 import re
 
 from sonic_py_common import device_info, logger
-from swsscommon.swsscommon import SonicV2Connector, ConfigDBConnector, SonicDBConfig
+from swsscommon.swsscommon import SonicV2Connector, ConfigDBConnector
 from minigraph import parse_xml
 from utilities_common.helper import update_config
+from utilities_common.general import load_db_config
 
 INIT_CFG_FILE = '/etc/sonic/init_cfg.json'
 MINIGRAPH_FILE = '/etc/sonic/minigraph.xml'
@@ -1256,10 +1257,7 @@ def main():
         socket_path = args.socket
         namespace = args.namespace
 
-        if args.namespace is not None:
-            SonicDBConfig.load_sonic_global_db_config(namespace=args.namespace)
-        else:
-            SonicDBConfig.initialize()
+        load_db_config()
 
         if socket_path:
             dbmgtr = DBMigrator(namespace, socket=socket_path)
diff --git a/tests/db_migrator_test.py b/tests/db_migrator_test.py
index 6ee4589a98..2e5544be8e 100644
--- a/tests/db_migrator_test.py
+++ b/tests/db_migrator_test.py
@@ -1,7 +1,8 @@
 import os
 import pytest
 import sys
-
+import argparse
+from unittest import mock
 from deepdiff import DeepDiff
 
 from swsscommon.swsscommon import SonicV2Connector
@@ -904,3 +905,18 @@ def test_golden_config_hostname(self):
         hostname = host.get('hostname', '')
         # hostname is from minigraph.xml
         assert hostname == 'SONiC-Dummy'
+
+class TestMain(object):
+    @classmethod
+    def setup_class(cls):
+        os.environ['UTILITIES_UNIT_TESTING'] = "2"
+
+    @classmethod
+    def teardown_class(cls):
+        os.environ['UTILITIES_UNIT_TESTING'] = "0"
+
+    @mock.patch('argparse.ArgumentParser.parse_args')
+    def test_init(self, mock_args):
+        mock_args.return_value=argparse.Namespace(namespace=None, operation='get_version', socket=None)
+        import db_migrator
+        db_migrator.main()