diff --git a/contracts/DecentralizedKV.sol b/contracts/DecentralizedKV.sol
index ce4d3a1..db82432 100644
--- a/contracts/DecentralizedKV.sol
+++ b/contracts/DecentralizedKV.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.28;
 
-import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
 import "./libraries/MerkleLib.sol";
 import "./libraries/BinaryRelated.sol";
 
@@ -9,7 +9,7 @@ import "./libraries/BinaryRelated.sol";
 /// @title DecentralizedKV
 /// @notice The DecentralizedKV is a top base contract for the EthStorage contract. It provides the
 ///         basic key-value store functionalities.
-contract DecentralizedKV is OwnableUpgradeable {
+contract DecentralizedKV is Ownable2StepUpgradeable {
     /// @notice Represents the metadata of the key-value .
     /// @custom:field kvIdx  Internal address seeking.
     /// @custom:field kvSize BLOB size.
diff --git a/contracts/test/DecentralizedKVTest.t.sol b/contracts/test/DecentralizedKVTest.t.sol
new file mode 100644
index 0000000..e1a0fb3
--- /dev/null
+++ b/contracts/test/DecentralizedKVTest.t.sol
@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.28;
+
+import "./TestDecentralizedKV.sol";
+import "forge-std/Test.sol";
+
+contract DecentralizedKVTest is Test {
+    uint256 constant MAX_KV_SIZE = 17;
+    uint256 constant STORAGE_COST = 10000000;
+    uint256 constant SHARD_SIZE_BITS = 19;
+    uint256 constant PREPAID_AMOUNT = 2 * STORAGE_COST;
+    TestDecentralizedKV decentralizedKV;
+
+    function setUp() public {
+        decentralizedKV = new TestDecentralizedKV(MAX_KV_SIZE, 0, STORAGE_COST, 340282366367469178095360967382638002176);
+        decentralizedKV.initialize(vm.addr(1));
+    }
+
+    function testTransferOwnership() public {
+        vm.expectRevert();
+        vm.prank(address(vm.addr(2)));
+        decentralizedKV.transferOwnership(vm.addr(2));
+
+        vm.prank(address(vm.addr(1)));
+        decentralizedKV.transferOwnership(vm.addr(2));
+        assertEq(decentralizedKV.owner(), vm.addr(1));
+
+        vm.prank(address(vm.addr(2)));
+        decentralizedKV.acceptOwnership();
+        assertEq(decentralizedKV.owner(), vm.addr(2));
+    }
+}