diff --git a/addresses.json b/addresses.json
index 89cf77f..995f1a6 100644
--- a/addresses.json
+++ b/addresses.json
@@ -11,7 +11,8 @@
     "GraphToken": "0x5c946740441C12510a167B447B7dE565C20b9E3C",
     "L1GraphTokenGateway": "0xc82fF7b51c3e593D709BA3dE1b3a0d233D1DEca1",
     "governor": "0xc6B8Dd6163a0d6E3Ff37Eb4306DeD0069FE4bF59",
-    "BillingConnector": "0x1c133bF1A059F682910d9bf81a09931742ec7311"
+    "BillingConnector": "0x1c133bF1A059F682910d9bf81a09931742ec7311",
+    "GelatoAutomate": "0x2A6C106ae13B558BB9E2Ec64Bd2f1f7BEFF3A5E0"
   },
   "137": {
     "Billing": "0x10829DB618E6F520Fa3A01c75bC6dDf8722fA9fE"
@@ -33,7 +34,7 @@
     "governor": "0xc6B8Dd6163a0d6E3Ff37Eb4306DeD0069FE4bF59",
     "Billing": "0x2da5DF23869B05b6641C91297b7025029112abA7",
     "BanxaWrapper": "0xDc4787B4c822a36F72A0084B02dA213A7291C35B",
-    "GelatoAutomate": "0xa5f9b728ecEB9A1F6FCC89dcc2eFd810bA4Dec41",
+    "GelatoAutomate": "0x2A6C106ae13B558BB9E2Ec64Bd2f1f7BEFF3A5E0",
     "RecurringPayments": "0x68becb01A921855D2b85B2952ca10f1801b518eD"
   }
 }
diff --git a/contracts/GelatoManager.sol b/contracts/GelatoManager.sol
index 49612c3..acb8980 100644
--- a/contracts/GelatoManager.sol
+++ b/contracts/GelatoManager.sol
@@ -23,20 +23,6 @@ contract GelatoManager is AutomateTaskCreator, Governed {
      */
     event MaxGasPriceSet(uint256 maxGasPrice);
 
-    /**
-     * @dev Emitted when funds are deposited into the Gelato treasury
-     * @param depositor User making the deposit
-     * @param amount Amount being deposited
-     */
-    event TreasuryFundsDeposited(address indexed depositor, uint256 amount);
-
-    /**
-     * @dev Emitted when funds are withdrawn from the Gelato treasury
-     * @param recepient Recepient receiving the funds
-     * @param amount Amount being withdrawn
-     */
-    event TreasuryFundsWithdrawn(address indexed recepient, uint256 amount);
-
     /**
      * @dev Emitted when a resolver task is created in Gelato Network
      * @param taskId The id of the task
@@ -74,31 +60,10 @@ contract GelatoManager is AutomateTaskCreator, Governed {
         address _automate,
         address _governor,
         uint256 _maxGasPrice
-    ) AutomateTaskCreator(_automate, _governor) Governed(_governor) {
+    ) AutomateTaskCreator(_automate) Governed(_governor) {
         _setMaxGasPrice(_maxGasPrice);
     }
 
-    /**
-     * @notice Deposit eth into the Gelato Network treasury
-     * This function is NOT meant to be used by the end user, only by the contract operator
-     * Any funds deposited via this call will be used to pay for Gelato Network tasks
-     */
-    function deposit() external payable {
-        if (msg.value == 0) revert InvalidDepositAmount();
-        taskTreasury.depositFunds{ value: msg.value }(address(this), ETH, msg.value);
-        emit TreasuryFundsDeposited(msg.sender, msg.value);
-    }
-
-    /**
-     * @notice Withdraw eth from the Gelato Network treasury
-     * @param recepient Recepient receiving the funds
-     * @param amount Amount being withdrawn
-     */
-    function withdraw(address recepient, uint256 amount) external onlyGovernor {
-        taskTreasury.withdrawFunds(payable(recepient), ETH, amount);
-        emit TreasuryFundsWithdrawn(recepient, amount);
-    }
-
     /**
      * @notice Sets the maximum gas price for task execution
      * @param newGasPrice The updated value for `maxGasPrice`
@@ -131,9 +96,13 @@ contract GelatoManager is AutomateTaskCreator, Governed {
         address execAddress,
         bytes memory execDataOrSelector
     ) internal returns (bytes32) {
-        ModuleData memory moduleData = ModuleData({ modules: new Module[](1), args: new bytes[](1) });
+        ModuleData memory moduleData = ModuleData({ modules: new Module[](2), args: new bytes[](2) });
+
         moduleData.modules[0] = Module.RESOLVER;
+        moduleData.modules[1] = Module.PROXY;
+
         moduleData.args[0] = _resolverModuleArg(resolverAddress, resolverData);
+        moduleData.args[1] = _proxyModuleArg();
 
         bytes32 taskId = _createTask(execAddress, execDataOrSelector, moduleData, address(0));
         emit ResolverTaskCreated(taskId);
diff --git a/contracts/gelato/AutomateReady.sol b/contracts/gelato/AutomateReady.sol
index d6c2ce2..b9d11fd 100644
--- a/contracts/gelato/AutomateReady.sol
+++ b/contracts/gelato/AutomateReady.sol
@@ -13,9 +13,8 @@ import "./Types.sol";
 abstract contract AutomateReady {
     IAutomate public immutable automate;
     address public immutable dedicatedMsgSender;
-    address private immutable _gelato;
+    address private immutable feeCollector;
     address internal constant ETH = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
-    address private constant OPS_PROXY_FACTORY = 0xC815dB16D4be6ddf2685C201937905aBf338F5D7;
 
     /**
      * @dev
@@ -33,8 +32,15 @@ abstract contract AutomateReady {
      */
     constructor(address _automate, address _taskCreator) {
         automate = IAutomate(_automate);
-        _gelato = IAutomate(_automate).gelato();
-        (dedicatedMsgSender, ) = IOpsProxyFactory(OPS_PROXY_FACTORY).getProxyOf(_taskCreator);
+        IGelato gelato = IGelato(IAutomate(_automate).gelato());
+
+        feeCollector = gelato.feeCollector();
+
+        address proxyModuleAddress = IAutomate(_automate).taskModuleAddresses(Module.PROXY);
+
+        address opsProxyFactoryAddress = IProxyModule(proxyModuleAddress).opsProxyFactory();
+
+        (dedicatedMsgSender, ) = IOpsProxyFactory(opsProxyFactoryAddress).getProxyOf(_taskCreator);
     }
 
     /**
@@ -45,10 +51,10 @@ abstract contract AutomateReady {
      */
     function _transfer(uint256 _fee, address _feeToken) internal {
         if (_feeToken == ETH) {
-            (bool success, ) = _gelato.call{ value: _fee }("");
+            (bool success, ) = feeCollector.call{ value: _fee }("");
             require(success, "_transfer: ETH transfer failed");
         } else {
-            SafeERC20.safeTransfer(IERC20(_feeToken), _gelato, _fee);
+            SafeERC20.safeTransfer(IERC20(_feeToken), feeCollector, _fee);
         }
     }
 
diff --git a/contracts/gelato/AutomateTaskCreator.sol b/contracts/gelato/AutomateTaskCreator.sol
index 2a87165..ccf3ed0 100644
--- a/contracts/gelato/AutomateTaskCreator.sol
+++ b/contracts/gelato/AutomateTaskCreator.sol
@@ -2,35 +2,35 @@
 pragma solidity ^0.8.14;
 
 import "./AutomateReady.sol";
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 
 /**
  * @dev Inherit this contract to allow your smart contract
  * to be a task creator and create tasks.
  */
+//solhint-disable const-name-snakecase
+//solhint-disable no-empty-blocks
 abstract contract AutomateTaskCreator is AutomateReady {
     using SafeERC20 for IERC20;
 
-    address public immutable fundsOwner;
-    ITaskTreasuryUpgradable public immutable taskTreasury;
+    IGelato1Balance public constant gelato1Balance = IGelato1Balance(0x7506C12a824d73D9b08564d5Afc22c949434755e);
 
-    constructor(address _automate, address _fundsOwner) AutomateReady(_automate, address(this)) {
-        fundsOwner = _fundsOwner;
-        taskTreasury = automate.taskTreasury();
-    }
-
-    /**
-     * @dev
-     * Withdraw funds from this contract's Gelato balance to fundsOwner.
-     */
-    function withdrawFunds(uint256 _amount, address _token) external {
-        require(msg.sender == fundsOwner, "Only funds owner can withdraw funds");
-
-        taskTreasury.withdrawFunds(payable(fundsOwner), _token, _amount);
-    }
+    constructor(address _automate) AutomateReady(_automate, address(this)) {}
 
-    function _depositFunds(uint256 _amount, address _token) internal {
-        uint256 ethValue = _token == ETH ? _amount : 0;
-        taskTreasury.depositFunds{ value: ethValue }(address(this), _token, _amount);
+    function _depositFunds1Balance(uint256 _amount, address _token, address _sponsor) internal {
+        if (_token == ETH) {
+            ///@dev Only deposit ETH on goerli for now.
+            require(block.chainid == 5, "Only deposit ETH on goerli");
+            gelato1Balance.depositNative{ value: _amount }(_sponsor);
+        } else {
+            ///@dev Only deposit USDC on polygon for now.
+            require(
+                block.chainid == 137 && _token == address(0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174),
+                "Only deposit USDC on polygon"
+            );
+            IERC20(_token).approve(address(gelato1Balance), _amount);
+            gelato1Balance.depositToken(_sponsor, _token, _amount);
+        }
     }
 
     function _createTask(
@@ -53,10 +53,6 @@ abstract contract AutomateTaskCreator is AutomateReady {
         return abi.encode(_resolverAddress, _resolverData);
     }
 
-    function _timeModuleArg(uint256 _startTime, uint256 _interval) internal pure returns (bytes memory) {
-        return abi.encode(uint128(_startTime), uint128(_interval));
-    }
-
     function _proxyModuleArg() internal pure returns (bytes memory) {
         return bytes("");
     }
@@ -64,4 +60,39 @@ abstract contract AutomateTaskCreator is AutomateReady {
     function _singleExecModuleArg() internal pure returns (bytes memory) {
         return bytes("");
     }
+
+    function _web3FunctionModuleArg(
+        string memory _web3FunctionHash,
+        bytes memory _web3FunctionArgsHex
+    ) internal pure returns (bytes memory) {
+        return abi.encode(_web3FunctionHash, _web3FunctionArgsHex);
+    }
+
+    function _timeTriggerModuleArg(uint128 _start, uint128 _interval) internal pure returns (bytes memory) {
+        bytes memory triggerConfig = abi.encode(_start, _interval);
+
+        return abi.encode(TriggerType.TIME, triggerConfig);
+    }
+
+    function _cronTriggerModuleArg(string memory _expression) internal pure returns (bytes memory) {
+        bytes memory triggerConfig = abi.encode(_expression);
+
+        return abi.encode(TriggerType.CRON, triggerConfig);
+    }
+
+    function _eventTriggerModuleArg(
+        address _address,
+        bytes32[][] memory _topics,
+        uint256 _blockConfirmations
+    ) internal pure returns (bytes memory) {
+        bytes memory triggerConfig = abi.encode(_address, _topics, _blockConfirmations);
+
+        return abi.encode(TriggerType.EVENT, triggerConfig);
+    }
+
+    function _blockTriggerModuleArg() internal pure returns (bytes memory) {
+        bytes memory triggerConfig = abi.encode(bytes(""));
+
+        return abi.encode(TriggerType.BLOCK, triggerConfig);
+    }
 }
diff --git a/contracts/gelato/Types.sol b/contracts/gelato/Types.sol
index 8b4b945..30b2fbc 100644
--- a/contracts/gelato/Types.sol
+++ b/contracts/gelato/Types.sol
@@ -3,9 +3,18 @@ pragma solidity ^0.8.12;
 
 enum Module {
     RESOLVER,
-    TIME,
+    DEPRECATED_TIME,
     PROXY,
-    SINGLE_EXEC
+    SINGLE_EXEC,
+    WEB3_FUNCTION,
+    TRIGGER
+}
+
+enum TriggerType {
+    TIME,
+    CRON,
+    EVENT,
+    BLOCK
 }
 
 struct ModuleData {
@@ -27,15 +36,23 @@ interface IAutomate {
 
     function gelato() external view returns (address payable);
 
-    function taskTreasury() external view returns (ITaskTreasuryUpgradable);
+    function taskModuleAddresses(Module) external view returns (address);
 }
 
-interface ITaskTreasuryUpgradable {
-    function depositFunds(address receiver, address token, uint256 amount) external payable;
-
-    function withdrawFunds(address payable receiver, address token, uint256 amount) external;
+interface IProxyModule {
+    function opsProxyFactory() external view returns (address);
 }
 
 interface IOpsProxyFactory {
     function getProxyOf(address account) external view returns (address, bool);
 }
+
+interface IGelato1Balance {
+    function depositNative(address _sponsor) external payable;
+
+    function depositToken(address _sponsor, address _token, uint256 _amount) external;
+}
+
+interface IGelato {
+    function feeCollector() external view returns (address);
+}
diff --git a/contracts/tests/TaskTreasuryMock.sol b/contracts/tests/TaskTreasuryMock.sol
deleted file mode 100644
index 19436af..0000000
--- a/contracts/tests/TaskTreasuryMock.sol
+++ /dev/null
@@ -1,16 +0,0 @@
-// SPDX-License-Identifier: GPL-3.0-or-later
-pragma solidity ^0.8.16;
-
-/**
- * @title Mock for the Gelato Task treasury contract
- */
-contract TaskTreasuryMock {
-    /// @dev Noop. Receive eth and do nothing.
-    function depositFunds(address receiver, address token, uint256 amount) external payable {}
-
-    /// @dev Send eth to receiver.
-    function withdrawFunds(address receiver, address token, uint256 amount) external {
-        (bool sent, ) = payable(receiver).call{ value: amount }("");
-        require(sent, "Failed to send Ether");
-    }
-}
diff --git a/contracts/tests/AutomateMock.sol b/contracts/tests/gelato/AutomateMock.sol
similarity index 52%
rename from contracts/tests/AutomateMock.sol
rename to contracts/tests/gelato/AutomateMock.sol
index 4a8bd90..9dccf6d 100644
--- a/contracts/tests/AutomateMock.sol
+++ b/contracts/tests/gelato/AutomateMock.sol
@@ -1,26 +1,24 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 pragma solidity ^0.8.16;
 
-import "../gelato/Types.sol";
+import "../../gelato/Types.sol";
 
 /**
  * @title Mock for the Gelato Automate contract
  */
 contract AutomateMock {
-    address internal gelatoAddr;
-    address internal treasuryAddr;
+    address public gelato;
+    address public feeCollector;
+    address public taskModuleAddress;
 
-    constructor(address _gelato, address _treasury) {
-        gelatoAddr = _gelato;
-        treasuryAddr = _treasury;
+    constructor(address _gelato, address _taskModuleAddress) {
+        gelato = _gelato;
+        taskModuleAddress = _taskModuleAddress;
+        feeCollector = address(0); // For this mock we don't care about the feeCollector
     }
 
-    function gelato() external view returns (address) {
-        return gelatoAddr;
-    }
-
-    function taskTreasury() external view returns (address) {
-        return treasuryAddr;
+    function taskModuleAddresses(Module) external view returns (address) {
+        return taskModuleAddress;
     }
 
     /// @dev This is NOT how Gelato computes task IDs
@@ -34,7 +32,7 @@ contract AutomateMock {
         return task;
     }
 
-    function cancelTask(bytes32 _taskId) external pure returns (bool) {
+    function cancelTask(bytes32) external pure returns (bool) {
         return true;
     }
 }
diff --git a/contracts/tests/gelato/GelatoMock.sol b/contracts/tests/gelato/GelatoMock.sol
new file mode 100644
index 0000000..30221b6
--- /dev/null
+++ b/contracts/tests/gelato/GelatoMock.sol
@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8.16;
+
+import { IGelato } from "../../gelato/Types.sol";
+
+/**
+ * @title Mock for the Gelato main contract
+ */
+contract GelatoMock is IGelato {
+    // We don't care about the feeCollector for this mock
+    function feeCollector() external view returns (address) {
+        return address(0);
+    }
+}
diff --git a/contracts/tests/OpsProxyFactoryMock.sol b/contracts/tests/gelato/OpsProxyFactoryMock.sol
similarity index 87%
rename from contracts/tests/OpsProxyFactoryMock.sol
rename to contracts/tests/gelato/OpsProxyFactoryMock.sol
index 69b777e..997a4f9 100644
--- a/contracts/tests/OpsProxyFactoryMock.sol
+++ b/contracts/tests/gelato/OpsProxyFactoryMock.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 pragma solidity ^0.8.16;
 
-import { IOpsProxyFactory } from "../gelato/Types.sol";
+import { IOpsProxyFactory } from "../../gelato/Types.sol";
 
 /**
  * @title Mock for the Gelato OpsProxyFactory contract
diff --git a/contracts/tests/gelato/ProxyModuleMock.sol b/contracts/tests/gelato/ProxyModuleMock.sol
new file mode 100644
index 0000000..c590ffd
--- /dev/null
+++ b/contracts/tests/gelato/ProxyModuleMock.sol
@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8.16;
+
+import { IProxyModule } from "../../gelato/Types.sol";
+
+/**
+ * @title Mock for the Gelato OpsProxyFactory contract
+ */
+contract ProxyModuleMock is IProxyModule {
+    address internal opsProxyFactoryAddress;
+
+    constructor(address _opsProxyFactoryAddress) {
+        opsProxyFactoryAddress = _opsProxyFactoryAddress;
+    }
+
+    function opsProxyFactory() external view returns (address) {
+        return opsProxyFactoryAddress;
+    }
+}
diff --git a/test/recurring-payments/contract.test.ts b/test/recurring-payments/contract.test.ts
index 16e762f..12d9efb 100644
--- a/test/recurring-payments/contract.test.ts
+++ b/test/recurring-payments/contract.test.ts
@@ -43,13 +43,13 @@ describe('RecurringPayments: Contract', () => {
 
   beforeEach(async function () {
     // eslint-disable-next-line @typescript-eslint/no-extra-semi
-    ;[me, governor, gelatoNetwork, user1] = await getAccounts()
+    ;[me, governor, user1] = await getAccounts()
 
     createData = ethers.utils.defaultAbiCoder.encode(['address', 'uint256'], [user1.address, oneHundred])
 
     token = await deployment.deployToken([tenBillion], me.signer, true)
 
-    automate = await deployMockGelatoNetwork(me.signer, gelatoNetwork.address)
+    automate = await deployMockGelatoNetwork(me.signer)
 
     // Deploy RecurringPayments contract
     recurringPayments = await deployment.deployRecurringPayments(
diff --git a/test/recurring-payments/manager.test.ts b/test/recurring-payments/manager.test.ts
index 2157842..55f9a47 100644
--- a/test/recurring-payments/manager.test.ts
+++ b/test/recurring-payments/manager.test.ts
@@ -53,10 +53,6 @@ describe('RecurringPayments: Gelato Manager', () => {
   })
 
   describe('constructor', function () {
-    it('should set gelato funds owner', async function () {
-      expect(await recurringPayments.fundsOwner()).to.eq(governor.address)
-    })
-
     it('should set gelato automate contract address', async function () {
       expect(await recurringPayments.automate()).to.eq(automate.address)
     })
@@ -102,57 +98,4 @@ describe('RecurringPayments: Gelato Manager', () => {
       await expect(recurringPayments.checkGasPrice({ gasPrice: initialMaxGasPrice })).not.to.be.reverted
     })
   })
-
-  describe('treasury', function () {
-    it('should revert when depositing zero eth', async function () {
-      const tx = recurringPayments.connect(governor.signer).deposit({ value: 0 })
-      await expect(tx).to.be.revertedWithCustomError(recurringPayments, 'InvalidDepositAmount')
-    })
-
-    it('should allow depositing eth', async function () {
-      const userBalanceBefore = await ethers.provider.getBalance(governor.address)
-      const tx = recurringPayments.connect(governor.signer).deposit({ value: ten })
-
-      await expect(tx).to.emit(recurringPayments, 'TreasuryFundsDeposited').withArgs(governor.address, ten)
-      const receipt = await (await tx).wait()
-
-      const userBalanceAfter = await ethers.provider.getBalance(governor.address)
-      expect(userBalanceAfter).to.eq(userBalanceBefore.sub(ten).sub(receipt.gasUsed.mul((await tx).gasPrice)))
-    })
-
-    it('should allow anyone to deposit', async function () {
-      const userBalanceBefore = await ethers.provider.getBalance(user1.address)
-      const tx = recurringPayments.connect(user1.signer).deposit({ value: ten })
-
-      await expect(tx).to.emit(recurringPayments, 'TreasuryFundsDeposited').withArgs(user1.address, ten)
-      const receipt = await (await tx).wait()
-
-      const userBalanceAfter = await ethers.provider.getBalance(user1.address)
-      expect(userBalanceAfter).to.eq(userBalanceBefore.sub(ten).sub(receipt.gasUsed.mul((await tx).gasPrice)))
-    })
-
-    it('should allow governor to withdraw', async function () {
-      const userBalanceBefore = await ethers.provider.getBalance(governor.address)
-
-      // Deposit
-      const tx = recurringPayments.connect(governor.signer).deposit({ value: ten })
-      await expect(tx).to.emit(recurringPayments, 'TreasuryFundsDeposited').withArgs(governor.address, ten)
-      const receipt = await (await tx).wait()
-
-      // Withdraw
-      const tx2 = recurringPayments.connect(governor.signer).withdraw(governor.address, ten)
-      await expect(tx2).to.emit(recurringPayments, 'TreasuryFundsWithdrawn').withArgs(governor.address, ten)
-      const receipt2 = await (await tx2).wait()
-
-      const userBalanceAfter = await ethers.provider.getBalance(governor.address)
-      expect(userBalanceAfter).to.eq(
-        userBalanceBefore.sub(receipt.gasUsed.mul((await tx).gasPrice)).sub(receipt2.gasUsed.mul((await tx2).gasPrice)),
-      )
-    })
-
-    it('should revert if unauthorized user attempts to withdraw', async function () {
-      const tx = recurringPayments.connect(user1.signer).withdraw(user1.address, ten)
-      await expect(tx).to.be.revertedWith('Only Governor can call')
-    })
-  })
 })
diff --git a/utils/deploy.ts b/utils/deploy.ts
index f47d0e7..dc040ac 100644
--- a/utils/deploy.ts
+++ b/utils/deploy.ts
@@ -9,9 +9,10 @@ import {
   AutomateMock,
   OpsProxyFactoryMock,
   PaymentMock,
+  ProxyModuleMock,
   SimplePaymentMock,
   Subscriptions,
-  TaskTreasuryMock,
+  GelatoMock,
 } from '../build/types'
 
 const hash = (input: string): string => utils.keccak256(`0x${input.replace(/^0x/, '')}`)
@@ -115,21 +116,30 @@ export async function deployAutomateMock(
 }
 
 // Pass the args in order to this func
-export async function deployTaskTreasuryMock(
+export async function deployOpsProxyFactoryMock(
   args: Array<string | BigNumber>,
   sender: Signer,
   disableLogging?: boolean,
-): Promise<TaskTreasuryMock> {
-  return deployContract(args, sender, 'TaskTreasuryMock', disableLogging) as unknown as Promise<TaskTreasuryMock>
+): Promise<OpsProxyFactoryMock> {
+  return deployContract(args, sender, 'OpsProxyFactoryMock', disableLogging) as unknown as Promise<OpsProxyFactoryMock>
 }
 
 // Pass the args in order to this func
-export async function deployOpsProxyFactoryMock(
+export async function deployProxyModuleMock(
   args: Array<string | BigNumber>,
   sender: Signer,
   disableLogging?: boolean,
-): Promise<OpsProxyFactoryMock> {
-  return deployContract(args, sender, 'OpsProxyFactoryMock', disableLogging) as unknown as Promise<OpsProxyFactoryMock>
+): Promise<ProxyModuleMock> {
+  return deployContract(args, sender, 'ProxyModuleMock', disableLogging) as unknown as Promise<ProxyModuleMock>
+}
+
+// Pass the args in order to this func
+export async function deployGelatoMock(
+  args: Array<string | BigNumber>,
+  sender: Signer,
+  disableLogging?: boolean,
+): Promise<GelatoMock> {
+  return deployContract(args, sender, 'GelatoMock', disableLogging) as unknown as Promise<GelatoMock>
 }
 
 // Pass the args in order to this func
diff --git a/utils/gelato.ts b/utils/gelato.ts
index 027a3f4..4d3af1f 100644
--- a/utils/gelato.ts
+++ b/utils/gelato.ts
@@ -1,19 +1,16 @@
-import hre from 'hardhat'
 import * as deployment from './deploy'
 import { Signer } from 'ethers'
-import * as OpsProxyFactoryArtifact from '../build/artifacts/contracts/tests/OpsProxyFactoryMock.sol/OpsProxyFactoryMock.json'
 
-export async function deployMockGelatoNetwork(deployer: Signer, gelato: string) {
-  // Deploy Gelato Task Treasury
-  const treasury = await deployment.deployTaskTreasuryMock([], deployer, true)
-
-  // Deploy Gelato Automate contract
-  const automate = await deployment.deployAutomateMock([gelato, treasury.address], deployer, true)
+export async function deployMockGelatoNetwork(deployer: Signer) {
+  // Deploy Gelato
+  const gelato = await deployment.deployGelatoMock([], deployer, true)
 
   // Deploy Gelato OpsProxyFactory
-  // The OpsProxyFactory address hardcoded in the Automate contract so we need to copy the bytecode into target address
-  const OPS_PROXY_FACTORY = '0xC815dB16D4be6ddf2685C201937905aBf338F5D7'
-  await hre.network.provider.send('hardhat_setCode', [OPS_PROXY_FACTORY, OpsProxyFactoryArtifact.deployedBytecode])
+  const opsProxyFactory = await deployment.deployOpsProxyFactoryMock([], deployer, true)
 
-  return automate
+  // Deploy Gelato Proxy module contract
+  const proxyModule = await deployment.deployProxyModuleMock([opsProxyFactory.address], deployer, true)
+
+  // Deploy Gelato Automate contract
+  return await deployment.deployAutomateMock([gelato.address, proxyModule.address], deployer, true)
 }