Fixes

contracts/MigrationZapper.sol

@@ -12,11 +12,17 @@ contract MigrationZapper {
     IMigrator public immutable migrator;
     IStaking public immutable ognStaking;
 
-    constructor(address _ogv, address _ogn, address _migrator, address _ognStaking) {
+    address public immutable governor;
+
+    error NotGovernor();
+
+    constructor(address _ogv, address _ogn, address _migrator, address _ognStaking, address _governor) {
         ogv = IMintableERC20(_ogv);
         ogn = IMintableERC20(_ogn);
         migrator = IMigrator(_migrator);
         ognStaking = IStaking(_ognStaking);
+
+        governor = _governor;
     }
 
     function initialize() external {
@@ -60,5 +66,23 @@ contract MigrationZapper {
             false,
             -1 // New stake
         );
+
+        // Transfer remaining OGN to the receiver
+        if (ognReceived > newStakeAmount) {
+            ogn.transfer(msg.sender, ognReceived - newStakeAmount);
+        }
+    }
+
+    /**
+     * Transfers any tokens sent by mistake out of the contract
+     * @param token Token address
+     * @param amount Amount of token to transfer
+     */
+    function transferTokens(address token, uint256 amount) external {
+        if (msg.sender != governor) {
+            revert NotGovernor();
+        }
+
+        IMintableERC20(token).transfer(governor, amount);
     }
 }

script/deploy/mainnet/012_MigrationZapperScript.sol

@@ -33,8 +33,9 @@ contract MigrationZapperScript is BaseMainnetScript {
         console.log("Deploy:");
         console.log("------------");
 
-        MigrationZapper zapper =
-            new MigrationZapper(Addresses.OGV, Addresses.OGN, deployedContracts["MIGRATOR"], deployedContracts["XOGN"]);
+        MigrationZapper zapper = new MigrationZapper(
+            Addresses.OGV, Addresses.OGN, deployedContracts["MIGRATOR"], deployedContracts["XOGN"], Addresses.TIMELOCK
+        );
         _recordDeploy("MIGRATION_ZAPPER", address(zapper));
 
         // Make sure Migrator can move OGV and OGN

tests/staking/ZapperForkTest.t.sol

@@ -84,4 +84,38 @@ contract ZapperForkTest is Test {
 
         vm.stopPrank();
     }
+
+    function testMigrateBalanceAndPartialStake() public {
+        vm.startPrank(ogvWhale);
+
+        uint256[] memory lockupIds = new uint256[](0);
+        uint256 stakeAmount = (50 ether * 0.09137 ether) / 1 ether;
+
+        uint256 ognBalanceBefore = ogn.balanceOf(ogvWhale);
+
+        zapper.migrate(100 ether, stakeAmount, 300 days);
+
+        // Should have it in a single OGN lockup
+        (uint128 amount, uint128 end, uint256 points) = xogn.lockups(ogvWhale, 0);
+        assertEq(amount, stakeAmount, "Balance not staked");
+
+        assertEq(ogn.balanceOf(ogvWhale), ognBalanceBefore + stakeAmount, "Less OGN received");
+
+        vm.stopPrank();
+    }
+
+    function testTransferTokens() public {
+        vm.startPrank(ogvWhale);
+        ogn.transfer(address(zapper), 100 ether);
+        vm.stopPrank();
+
+        vm.startPrank(Addresses.TIMELOCK);
+        zapper.transferTokens(address(ogn), 100 ether);
+        vm.stopPrank();
+
+        vm.startPrank(ogvWhale);
+        vm.expectRevert(bytes4(keccak256("NotGovernor()")));
+        zapper.transferTokens(address(ogn), 100 ether);
+        vm.stopPrank();
+    }
 }