diff --git a/contracts/tokens/SecurityToken.sol b/contracts/tokens/SecurityToken.sol
index fa915e5e1..d8b2ddfe4 100644
--- a/contracts/tokens/SecurityToken.sol
+++ b/contracts/tokens/SecurityToken.sol
@@ -526,21 +526,18 @@ contract SecurityToken is ERC20, ReentrancyGuard, SecurityTokenStorage, IERC1594
         // require(balanceOfByPartition(_partition, msg.sender) >= _value);
         // NB - Above condition will be automatically checked using the executeTransfer() function execution.
         // NB - passing `_additionalBalance` value is 0 because accessing the balance before transfer
-        uint256 balanceBeforeTransferLocked = _balanceOfByPartition(_partition, _to, 0);
+        uint256 lockedBalanceBeforeTransfer = _balanceOfByPartition(LOCKED, _to, 0);
         _transferWithData(_from, _to, _value, _data);
         // NB - passing `_additonalBalance` valie is 0 because balance of `_to` was updated in the transfer call
-        uint256 balanceAfterTransferLocked = _balanceOfByPartition(_partition, _to, 0);
-        toPartition =  _returnPartition(balanceBeforeTransferLocked, balanceAfterTransferLocked, _value);
+        uint256 lockedBalanceAfterTransfer = _balanceOfByPartition(LOCKED, _to, 0);
+        toPartition =  _returnPartition(lockedBalanceBeforeTransfer, lockedBalanceAfterTransfer, _value);
         emit TransferByPartition(_partition, _operator, _from, _to, _value, _data, _operatorData);
     }
 
     function _returnPartition(uint256 _beforeBalance, uint256 _afterBalance, uint256 _value) internal pure returns(bytes32 toPartition) {
         // return LOCKED only when the transaction `_value` should be equal to the change in the LOCKED partition
         // balance otherwise return UNLOCKED
-        if (_afterBalance.sub(_beforeBalance) == _value)
-            toPartition = LOCKED;
-        // Returning the same partition UNLOCKED
-        toPartition = UNLOCKED;
+        toPartition = _afterBalance.sub(_beforeBalance) == _value ? LOCKED : UNLOCKED; // Returning the same partition UNLOCKED
     }
 
     ///////////////////////
@@ -974,8 +971,8 @@ contract SecurityToken is ERC20, ReentrancyGuard, SecurityTokenStorage, IERC1594
             bool success;
             (success, esc, appStatusCode) = _canTransfer(_from, _to, _value, _data);
             if (success) {
-                uint256 beforeBalance = _balanceOfByPartition(_partition, _to, 0);
-                uint256 afterbalance = _balanceOfByPartition(_partition, _to, _value);
+                uint256 beforeBalance = _balanceOfByPartition(LOCKED, _to, 0);
+                uint256 afterbalance = _balanceOfByPartition(LOCKED, _to, _value);
                 toPartition = _returnPartition(beforeBalance, afterbalance, _value);
             }
             return (esc, appStatusCode, toPartition);
diff --git a/test/o_security_token.js b/test/o_security_token.js
index 71afe9be5..7e4291f32 100644
--- a/test/o_security_token.js
+++ b/test/o_security_token.js
@@ -9,7 +9,8 @@ import {
     deployGPMAndVerifyed,
     deployCappedSTOAndVerifyed,
     deployMockRedemptionAndVerifyed,
-    deployMockWrongTypeRedemptionAndVerifyed
+    deployMockWrongTypeRedemptionAndVerifyed,
+    deployLockUpTMAndVerified
 } from "./helpers/createInstances";
 
 const MockSecurityTokenLogic = artifacts.require("./MockSecurityTokenLogic.sol");
@@ -20,6 +21,7 @@ const SecurityToken = artifacts.require("./SecurityToken.sol");
 const GeneralTransferManager = artifacts.require("./GeneralTransferManager");
 const GeneralPermissionManager = artifacts.require("./GeneralPermissionManager");
 const MockRedemptionManager = artifacts.require("./MockRedemptionManager.sol");
+const LockUpTransferManager = artifacts.require("./LockUpTransferManager.sol");
 const STGetter = artifacts.require("./STGetter.sol");
 
 const Web3 = require("web3");
@@ -61,6 +63,8 @@ contract("SecurityToken", async (accounts) => {
 
     // Contract Instance Declaration
     let I_GeneralPermissionManagerFactory;
+    let I_LockUpTransferManagerFactory;
+    let I_LockUpTransferManager;
     let I_SecurityTokenRegistryProxy;
     let I_GeneralTransferManagerFactory;
     let I_GeneralPermissionManager;
@@ -162,6 +166,8 @@ contract("SecurityToken", async (accounts) => {
         [I_GeneralPermissionManagerFactory] = await deployGPMAndVerifyed(account_polymath, I_MRProxied, 0);
         // STEP 3: Deploy the CappedSTOFactory
         [I_CappedSTOFactory] = await deployCappedSTOAndVerifyed(account_polymath, I_MRProxied, cappedSTOSetupCost);
+        // STEP 4(c): Deploy the LockUpVolumeRestrictionTMFactory
+        [I_LockUpTransferManagerFactory] = await deployLockUpTMAndVerified(account_polymath, I_MRProxied, 0);
 
         // Printing all the contract addresses
         console.log(`
@@ -2172,6 +2178,57 @@ contract("SecurityToken", async (accounts) => {
                 assert.equal(web3.utils.toUtf8(allDocs[0]), "doc4");
             });
         });
+
+        describe("Test cases for the returnPartition", async() => {
+            // It will work once the balanceOfByPartition function fixed added
+            it.skip("Should add the lockup Transfer manager and create a lockup for investor 1", async() => {
+
+                console.log(web3.utils.fromWei(await I_SecurityToken.balanceOf.call(account_investor1)));
+                console.log(web3.utils.fromWei(await I_SecurityToken.balanceOfByPartition.call(web3.utils.toHex("UNLOCKED"),account_investor1)));
+                console.log(web3.utils.fromWei(await I_SecurityToken.balanceOf.call(account_investor2)));
+
+                const tx = await I_SecurityToken.addModule(I_LockUpTransferManagerFactory.address, "0x", 0, 0, false, { from: token_owner });
+                assert.equal(tx.logs[2].args._types[0].toString(), transferManagerKey, "LockUpVolumeRestrictionTMFactory doesn't get deployed");
+                assert.equal(
+                    web3.utils.toAscii(tx.logs[2].args._name)
+                    .replace(/\u0000/g, ''),
+                    "LockUpTransferManager",
+                    "LockUpTransferManager module was not added"
+                );
+                I_LockUpTransferManager = await LockUpTransferManager.at(tx.logs[2].args._module);
+                let currentTime = new BN(await latestTime());
+                await I_LockUpTransferManager.addNewLockUpToUser(
+                    account_investor2,
+                    new BN(web3.utils.toWei("1000")),
+                    currentTime.add(new BN(duration.seconds(1))),
+                    new BN(duration.seconds(400000)),
+                    new BN(duration.seconds(100000)),
+                    web3.utils.fromAscii("a_lockup"),
+                    {
+                        from: token_owner
+                    }
+                );
+
+                // transfer balance of Unlocked partition of invesotor 1 to 2
+                await increaseTime(10);
+
+                console.log(`UNLOCKED balance - ${web3.utils.fromWei(await I_SecurityToken.balanceOfByPartition.call(web3.utils.toHex("UNLOCKED"),account_investor2))}`);
+                console.log(`Locked Balance - ${web3.utils.fromWei(await I_SecurityToken.balanceOfByPartition.call(web3.utils.toHex("LOCKED"),account_investor2))}`);
+
+                let partition = await I_SecurityToken.transferByPartition.call(
+                    web3.utils.toHex("UNLOCKED"),
+                    account_investor2,
+                    new BN(web3.utils.toWei("500")),
+                    "0x0",
+                    {
+                        from: account_investor1
+                    }
+                );
+                console.log(`UNLOCKED balance - ${web3.utils.fromWei(await I_SecurityToken.balanceOfByPartition.call(web3.utils.toHex("UNLOCKED"),account_investor2))}`);
+                console.log(`Locked Balance - ${web3.utils.fromWei(await I_SecurityToken.balanceOfByPartition.call(web3.utils.toHex("LOCKED"),account_investor2))}`);
+                assert.equal(web3.utils.hexToUtf8(partition), "LOCKED");
+            });
+        })
     })
     });