Merge branch 'master' into smart_switch_db

Signed-off-by: Ze Gan <[email protected]>

Author: Pterosaur

common/configdb.h

@@ -59,6 +59,13 @@ class ConfigDBConnector_Native : public SonicV2Connector_Native
             self.handlers = {}
             self.fire_init_data = {}
 
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_value, exc_tb):
+            self.close()
+            pass
+
         @property
         def KEY_SEPARATOR(self):
             return self.getKeySeparator()

common/dbinterface.cpp

@@ -32,6 +32,11 @@ void DBInterface::close(const std::string& dbName)
     m_redisClient.erase(dbName);
 }
 
+void DBInterface::close()
+{
+    m_redisClient.clear();
+}
+
 int64_t DBInterface::del(const string& dbName, const std::string& key, bool blocking)
 {
     auto innerfunc = [&]

common/dbinterface.h

@@ -34,6 +34,7 @@ class DBInterface
 public:
     void connect(int dbId, const std::string& dbName, bool retry = true);
     void close(const std::string& dbName);
+    void close();
     int64_t del(const std::string& dbName, const std::string& key, bool blocking = false);
     // Delete all keys which match %pattern from DB
     void delete_all_by_pattern(const std::string& dbName, const std::string& pattern);

common/sonicv2connector.cpp

@@ -34,6 +34,11 @@ void SonicV2Connector_Native::close(const std::string& db_name)
     m_dbintf.close(db_name);
 }
 
+void SonicV2Connector_Native::close()
+{
+    m_dbintf.close();
+}
+
 std::vector<std::string> SonicV2Connector_Native::get_db_list()
 {
     return SonicDBConfig::getDbList(m_netns);

common/sonicv2connector.h

@@ -20,6 +20,8 @@ class SonicV2Connector_Native
 
     void close(const std::string& db_name);
 
+    void close();
+
     std::vector<std::string> get_db_list();
 
     int get_dbid(const std::string& db_name);

tests/redis_ut.cpp

@@ -326,6 +326,7 @@ TEST(DBConnector, DBInterface)
     db.set("TEST_DB", "key0", "field1", "value2");
     auto fvs = db.get_all("TEST_DB", "key0");
     auto rc = fvs.find("field1");
+    db.close();
     EXPECT_NE(rc, fvs.end());
     EXPECT_EQ(rc->second, "value2");
 }

tests/test_redis_ut.py

@@ -7,6 +7,15 @@
 from swsscommon import swsscommon
 from swsscommon.swsscommon import ConfigDBPipeConnector, DBInterface, SonicV2Connector, SonicDBConfig, ConfigDBConnector, SonicDBConfig, transpose_pops, SonicDBKey
 import json
+import gc
+
+import sys
+if sys.version_info.major == 3:
+    from unittest import mock
+else:
+    # Expect the 'mock' package for python 2
+    # https://pypi.python.org/pypi/mock
+    import mock
 
 def test_ProducerTable():
     db = swsscommon.DBConnector("APPL_DB", 0, True)
@@ -804,6 +813,24 @@ def test_ConfigDBConnector():
     assert len(allconfig) == 0
 
 
+@mock.patch("swsscommon.swsscommon.ConfigDBConnector.close")
+def test_ConfigDBConnector_with_statement(self):
+    # test ConfigDBConnector support 'with' statement
+    with ConfigDBConnector() as config_db:
+        assert config_db.db_name == ""
+        assert config_db.TABLE_NAME_SEPARATOR == "|"
+        config_db.connect(wait_for_init=False)
+        assert config_db.db_name == "CONFIG_DB"
+        assert config_db.TABLE_NAME_SEPARATOR == "|"
+        config_db.get_redis_client(config_db.CONFIG_DB).flushdb()
+        config_db.set_entry("TEST_PORT", "Ethernet111", {"alias": "etp1x"})
+        allconfig = config_db.get_config()
+        assert allconfig["TEST_PORT"]["Ethernet111"]["alias"] == "etp1x"
+
+    # check close() method called by with statement
+    ConfigDBConnector.close.assert_called_once_with()
+
+
 def test_SmartSwitchDBConnector():
     test_dir = os.path.dirname(os.path.abspath(__file__))
     global_db_config = os.path.join(test_dir, 'redis_multi_db_ut_config', 'database_global.json')
@@ -821,3 +848,4 @@ def test_SmartSwitchDBConnector():
     assert "dputest2" in keys
     assert tbl.get("dputest1")[1][0] == ("dashfield1", "dashvalue1")
     assert tbl.get("dputest2")[1][1] == ("dashfield2", "dashvalue2")
+