diff --git a/src/libraries/MapWithTimeData.sol b/src/libraries/MapWithTimeData.sol
index 934be93..9b8540c 100644
--- a/src/libraries/MapWithTimeData.sol
+++ b/src/libraries/MapWithTimeData.sol
@@ -1,4 +1,17 @@
-// SPDX-License-Identifier: MIT
+//SPDX-License-Identifier: GPL-3.0-or-later
+
+// Copyright (C) Moondance Labs Ltd.
+// This file is part of Tanssi.
+// Tanssi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// Tanssi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License
+// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>MIT
 pragma solidity ^0.8.25;
 
 import {Checkpoints} from "@openzeppelin/contracts/utils/structs/Checkpoints.sol";
diff --git a/src/libraries/SimpleKeyRegistry32.sol b/src/libraries/SimpleKeyRegistry32.sol
index 5255d8f..0ca2e0a 100644
--- a/src/libraries/SimpleKeyRegistry32.sol
+++ b/src/libraries/SimpleKeyRegistry32.sol
@@ -1,4 +1,17 @@
-// SPDX-License-Identifier: MIT
+//SPDX-License-Identifier: GPL-3.0-or-later
+
+// Copyright (C) Moondance Labs Ltd.
+// This file is part of Tanssi.
+// Tanssi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// Tanssi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License
+// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
 pragma solidity 0.8.25;
 
 import {Checkpoints} from "@openzeppelin/contracts/utils/structs/Checkpoints.sol";
diff --git a/src/middleware/Middleware.sol b/src/middleware/Middleware.sol
index 05f1818..a3e4ed9 100644
--- a/src/middleware/Middleware.sol
+++ b/src/middleware/Middleware.sol
@@ -1,7 +1,19 @@
-// SPDX-License-Identifier: MIT
+//SPDX-License-Identifier: GPL-3.0-or-later
+
+// Copyright (C) Moondance Labs Ltd.
+// This file is part of Tanssi.
+// Tanssi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// Tanssi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License
+// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
 pragma solidity 0.8.25;
 
-import {console2} from "forge-std/console2.sol";
 import {Time} from "@openzeppelin/contracts/utils/types/Time.sol";
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 import {EnumerableMap} from "@openzeppelin/contracts/utils/structs/EnumerableMap.sol";
@@ -25,55 +37,59 @@ contract Middleware is SimpleKeyRegistry32, Ownable {
     using MapWithTimeData for EnumerableMap.AddressToUintMap;
     using Subnetwork for address;
 
-    error NotOperator();
-    error NotVault();
-
-    error OperatorNotOptedIn();
-    error OperatorNotRegistred();
-    error OperarorGracePeriodNotPassed();
-    error OperatorAlreadyRegistred();
-
-    error VaultAlreadyRegistered();
-    error VaultEpochTooShort();
-    error VaultGracePeriodNotPassed();
-
-    error InvalidSubnetworksCnt();
-
-    error TooOldEpoch();
-    error InvalidEpoch();
-
-    error SlashingWindowTooShort();
-    error TooBigSlashAmount();
-    error UnknownSlasherType();
+    error Middleware__NotOperator();
+    error Middleware__NotVault();
+    error Middleware__OperatorNotOptedIn();
+    error Middleware__OperatorNotRegistred();
+    error Middleware__OperarorGracePeriodNotPassed();
+    error Middleware__OperatorAlreadyRegistred();
+    error Middleware__VaultAlreadyRegistered();
+    error Middleware__VaultEpochTooShort();
+    error Middleware__VaultGracePeriodNotPassed();
+    error Middleware__InvalidSubnetworksCnt();
+    error Middleware__TooOldEpoch();
+    error Middleware__InvalidEpoch();
+    error Middleware__SlashingWindowTooShort();
+    error Middleware__TooBigSlashAmount();
+    error Middleware__UnknownSlasherType();
 
     struct ValidatorData {
         uint256 stake;
         bytes32 key;
     }
 
-    address public immutable NETWORK;
-    address public immutable OPERATOR_REGISTRY;
-    address public immutable VAULT_REGISTRY;
-    address public immutable OPERATOR_NET_OPTIN;
-    address public immutable OWNER;
-    uint48 public immutable EPOCH_DURATION;
-    uint48 public immutable SLASHING_WINDOW;
-    uint48 public immutable START_TIME;
+    struct SlashParams {
+        uint48 epochStartTs;
+        address vault;
+        address operator;
+        uint256 totalOperatorStake;
+        uint256 slashAmount;
+    }
+
+    address public immutable i_network;
+    address public immutable i_operatorRegistry;
+    address public immutable i_vaultRegistry;
+    address public immutable i_operatorNetworkOptin;
+    address public immutable i_owner;
+    uint48 public immutable i_epochDuration;
+    uint48 public immutable i_slashingWindow;
+
+    uint48 public immutable i_startTime;
 
     uint48 private constant INSTANT_SLASHER_TYPE = 0;
     uint48 private constant VETO_SLASHER_TYPE = 1;
 
-    uint256 public subnetworksCnt;
-    mapping(uint48 => uint256) public totalStakeCache;
-    mapping(uint48 => bool) public totalStakeCached;
-    mapping(uint48 => mapping(address => uint256)) public operatorStakeCache;
-    EnumerableMap.AddressToUintMap private operators;
-    EnumerableMap.AddressToUintMap private vaults;
+    uint256 public s_subnetworksCount;
+    mapping(uint48 => uint256) public s_totalStakeCache;
+    mapping(uint48 => bool) public s_totalStakeCached;
+    mapping(uint48 => mapping(address => uint256)) public s_operatorStakeCache;
+    EnumerableMap.AddressToUintMap private s_operators;
+    EnumerableMap.AddressToUintMap private s_vaults;
 
     modifier updateStakeCache(
         uint48 epoch
     ) {
-        if (!totalStakeCached[epoch]) {
+        if (!s_totalStakeCached[epoch]) {
             calcAndCacheStakes(epoch);
         }
         _;
@@ -89,96 +105,112 @@ contract Middleware is SimpleKeyRegistry32, Ownable {
         uint48 _slashingWindow
     ) SimpleKeyRegistry32() Ownable(_owner) {
         if (_slashingWindow < _epochDuration) {
-            revert SlashingWindowTooShort();
+            revert Middleware__SlashingWindowTooShort();
         }
-        START_TIME = Time.timestamp();
-        EPOCH_DURATION = _epochDuration;
-        NETWORK = _network;
-        OWNER = _owner;
-        OPERATOR_REGISTRY = _operatorRegistry;
-        VAULT_REGISTRY = _vaultRegistry;
-        OPERATOR_NET_OPTIN = _operatorNetOptin;
-        SLASHING_WINDOW = _slashingWindow;
-
-        subnetworksCnt = 1;
-    }
-
-    function getEpochStartTs(
-        uint48 epoch
-    ) public view returns (uint48 timestamp) {
-        return START_TIME + epoch * EPOCH_DURATION;
-    }
-
-    function getEpochAtTs(
-        uint48 timestamp
-    ) public view returns (uint48 epoch) {
-        return (timestamp - START_TIME) / EPOCH_DURATION;
-    }
-
-    function getCurrentEpoch() public view returns (uint48 epoch) {
-        return getEpochAtTs(Time.timestamp());
+        i_startTime = Time.timestamp();
+        i_epochDuration = _epochDuration;
+        i_network = _network;
+        i_owner = _owner;
+        i_operatorRegistry = _operatorRegistry;
+        i_vaultRegistry = _vaultRegistry;
+        i_operatorNetworkOptin = _operatorNetOptin;
+        i_slashingWindow = _slashingWindow;
+
+        s_subnetworksCount = 1;
     }
 
+    /**
+     * @notice Registers a new operator with a key
+     * @dev Only the owner can call this function
+     * @param operator The operator's address
+     * @param key The operator's key
+     */
     function registerOperator(address operator, bytes32 key) external onlyOwner {
-        if (operators.contains(operator)) {
-            revert OperatorAlreadyRegistred();
+        if (s_operators.contains(operator)) {
+            revert Middleware__OperatorAlreadyRegistred();
         }
 
-        if (!IRegistry(OPERATOR_REGISTRY).isEntity(operator)) {
-            revert NotOperator();
+        if (!IRegistry(i_operatorRegistry).isEntity(operator)) {
+            revert Middleware__NotOperator();
         }
 
-        if (!IOptInService(OPERATOR_NET_OPTIN).isOptedIn(operator, NETWORK)) {
-            revert OperatorNotOptedIn();
+        if (!IOptInService(i_operatorNetworkOptin).isOptedIn(operator, i_network)) {
+            revert Middleware__OperatorNotOptedIn();
         }
 
         updateKey(operator, key);
 
-        operators.add(operator);
-        operators.enable(operator);
+        s_operators.add(operator);
+        s_operators.enable(operator);
     }
 
+    /**
+     * @notice Updates an existing operator's key
+     * @dev Only the owner can call this function
+     * @param operator The operator's address
+     * @param key The new key
+     */
     function updateOperatorKey(address operator, bytes32 key) external onlyOwner {
-        if (!operators.contains(operator)) {
-            revert OperatorNotRegistred();
+        if (!s_operators.contains(operator)) {
+            revert Middleware__OperatorNotRegistred();
         }
 
         updateKey(operator, key);
     }
 
+    /**
+     * @notice Pauses an operator
+     * @dev Only the owner can call this function
+     * @param operator The operator to pause
+     */
     function pauseOperator(
         address operator
     ) external onlyOwner {
-        operators.disable(operator);
+        s_operators.disable(operator);
     }
 
+    /**
+     * @notice Re-enables a paused operator
+     * @dev Only the owner can call this function
+     * @param operator The operator to unpause
+     */
     function unpauseOperator(
         address operator
     ) external onlyOwner {
-        operators.enable(operator);
+        s_operators.enable(operator);
     }
 
+    /**
+     * @notice Removes an operator after grace period
+     * @dev Only the owner can call this function
+     * @param operator The operator to unregister
+     */
     function unregisterOperator(
         address operator
     ) external onlyOwner {
-        (, uint48 disabledTime) = operators.getTimes(operator);
+        (, uint48 disabledTime) = s_operators.getTimes(operator);
 
-        if (disabledTime == 0 || disabledTime + SLASHING_WINDOW > Time.timestamp()) {
-            revert OperarorGracePeriodNotPassed();
+        if (disabledTime == 0 || disabledTime + i_slashingWindow > Time.timestamp()) {
+            revert Middleware__OperarorGracePeriodNotPassed();
         }
 
-        operators.remove(operator);
+        s_operators.remove(operator);
     }
 
+    /**
+     * @notice Registers a new vault
+     * @dev Only the owner can call this function
+     * @param vault The vault address to register
+     */
     function registerVault(
         address vault
     ) external onlyOwner {
-        if (vaults.contains(vault)) {
-            revert VaultAlreadyRegistered();
+        if (s_vaults.contains(vault)) {
+            revert Middleware__VaultAlreadyRegistered();
         }
 
-        if (!IRegistry(VAULT_REGISTRY).isEntity(vault)) {
-            revert NotVault();
+        if (!IRegistry(i_vaultRegistry).isEntity(vault)) {
+            revert Middleware__NotVault();
         }
 
         uint48 vaultEpoch = IVault(vault).epochDuration();
@@ -188,134 +220,120 @@ contract Middleware is SimpleKeyRegistry32, Ownable {
             vaultEpoch -= IVetoSlasher(slasher).vetoDuration();
         }
 
-        if (vaultEpoch < SLASHING_WINDOW) {
-            revert VaultEpochTooShort();
+        if (vaultEpoch < i_slashingWindow) {
+            revert Middleware__VaultEpochTooShort();
         }
 
-        vaults.add(vault);
-        vaults.enable(vault);
-    }
-
-    function isVaultRegistered(
-        address vault
-    ) external view returns (bool) {
-        return vaults.contains(vault);
+        s_vaults.add(vault);
+        s_vaults.enable(vault);
     }
 
+    /**
+     * @notice Pauses a vault
+     * @dev Only the owner can call this function
+     * @param vault The vault to pause
+     */
     function pauseVault(
         address vault
     ) external onlyOwner {
-        vaults.disable(vault);
+        s_vaults.disable(vault);
     }
 
+    /**
+     * @notice Re-enables a paused vault
+     * @dev Only the owner can call this function
+     * @param vault The vault to unpause
+     */
     function unpauseVault(
         address vault
     ) external onlyOwner {
-        vaults.enable(vault);
+        s_vaults.enable(vault);
     }
 
+    /**
+     * @notice Removes a vault after grace period
+     * @dev Only the owner can call this function
+     * @param vault The vault to unregister
+     */
     function unregisterVault(
         address vault
     ) external onlyOwner {
-        (, uint48 disabledTime) = vaults.getTimes(vault);
+        (, uint48 disabledTime) = s_vaults.getTimes(vault);
 
-        if (disabledTime == 0 || disabledTime + SLASHING_WINDOW > Time.timestamp()) {
-            revert VaultGracePeriodNotPassed();
+        if (disabledTime == 0 || disabledTime + i_slashingWindow > Time.timestamp()) {
+            revert Middleware__VaultGracePeriodNotPassed();
         }
 
-        vaults.remove(vault);
+        s_vaults.remove(vault);
     }
 
-    function setSubnetworksCnt(
-        uint256 _subnetworksCnt
+    /**
+     * @notice Updates the number of subnetworks
+     * @dev Only the owner can call this function
+     * @param _subnetworksCount New subnetwork count
+     */
+    function setSubnetworksCount(
+        uint256 _subnetworksCount
     ) external onlyOwner {
-        if (subnetworksCnt >= _subnetworksCnt) {
-            revert InvalidSubnetworksCnt();
+        if (s_subnetworksCount >= _subnetworksCount) {
+            revert Middleware__InvalidSubnetworksCnt();
         }
 
-        subnetworksCnt = _subnetworksCnt;
+        s_subnetworksCount = _subnetworksCount;
     }
 
-    function getOperatorStake(address operator, uint48 epoch) public view returns (uint256 stake) {
-        if (totalStakeCached[epoch]) {
-            return operatorStakeCache[epoch][operator];
-        }
-
+    // function submission(bytes memory payload, bytes32[] memory signatures) public updateStakeCache(getCurrentEpoch()) {
+    //     // validate signatures
+    //     // validate payload
+    //     // process payload
+    // }
+
+    /**
+     * @notice Calculates and caches stakes for an epoch
+     * @param epoch The epoch to calculate for
+     * @return totalStake The total stake amount
+     */
+    function calcAndCacheStakes(
+        uint48 epoch
+    ) public returns (uint256 totalStake) {
         uint48 epochStartTs = getEpochStartTs(epoch);
-        for (uint256 i; i < vaults.length(); ++i) {
-            (address vault, uint48 enabledTime, uint48 disabledTime) = vaults.atWithTimes(i);
-
-            // just skip the vault if it was enabled after the target epoch or not enabled
-            if (!_wasActiveAt(enabledTime, disabledTime, epochStartTs)) {
-                continue;
-            }
 
-            for (uint96 j = 0; j < subnetworksCnt; ++j) {
-                stake += IBaseDelegator(IVault(vault).delegator()).stakeAt(
-                    NETWORK.subnetwork(j), operator, epochStartTs, new bytes(0)
-                );
-            }
+        // for epoch older than SLASHING_WINDOW total stake can be invalidated (use cache)
+        if (epochStartTs < Time.timestamp() - i_slashingWindow) {
+            revert Middleware__TooOldEpoch();
         }
 
-        return stake;
-    }
-
-    function getTotalStake(
-        uint48 epoch
-    ) public view returns (uint256) {
-        if (totalStakeCached[epoch]) {
-            return totalStakeCache[epoch];
+        if (epochStartTs > Time.timestamp()) {
+            revert Middleware__InvalidEpoch();
         }
-        return _calcTotalStake(epoch);
-    }
-
-    function getValidatorSet(
-        uint48 epoch
-    ) public view returns (ValidatorData[] memory validatorsData) {
-        uint48 epochStartTs = getEpochStartTs(epoch);
-
-        validatorsData = new ValidatorData[](operators.length());
-        uint256 valIdx = 0;
 
-        for (uint256 i; i < operators.length(); ++i) {
-            (address operator, uint48 enabledTime, uint48 disabledTime) = operators.atWithTimes(i);
+        for (uint256 i; i < s_operators.length(); ++i) {
+            (address operator, uint48 enabledTime, uint48 disabledTime) = s_operators.atWithTimes(i);
 
             // just skip operator if it was added after the target epoch or paused
             if (!_wasActiveAt(enabledTime, disabledTime, epochStartTs)) {
                 continue;
             }
 
-            bytes32 key = getOperatorKeyAt(operator, epochStartTs);
-            if (key == bytes32(0)) {
-                continue;
-            }
-
-            uint256 stake = getOperatorStake(operator, epoch);
-
-            validatorsData[valIdx++] = ValidatorData(stake, key);
-        }
+            uint256 operatorStake = getOperatorStake(operator, epoch);
+            s_operatorStakeCache[epoch][operator] = operatorStake;
 
-        // shrink array to skip unused slots
-        /// @solidity memory-safe-assembly
-        assembly {
-            mstore(validatorsData, valIdx)
+            totalStake += operatorStake;
         }
-    }
-
-    function submission(bytes memory payload, bytes32[] memory signatures) public updateStakeCache(getCurrentEpoch()) {
-        // validate signatures
-        // validate payload
-        // process payload
-    }
 
-    struct SlashParams {
-        uint48 epochStartTs;
-        address vault;
-        address operator;
-        uint256 totalOperatorStake;
-        uint256 slashAmount;
+        s_totalStakeCached[epoch] = true;
+        s_totalStakeCache[epoch] = totalStake;
     }
 
+    /**
+     * @notice Slashes an operator's stake
+     * @dev Only the owner can call this function
+     * @dev This function first updates the stake cache for the target epoch
+     * @param epoch The epoch number
+     * @param operator The operator to slash
+     * @param amount Amount to slash
+     */
+    //INFO: this function can be made external. To check if it is possible to make it external
     function slash(uint48 epoch, address operator, uint256 amount) public onlyOwner updateStakeCache(epoch) {
         SlashParams memory params;
         params.epochStartTs = getEpochStartTs(epoch);
@@ -325,20 +343,14 @@ contract Middleware is SimpleKeyRegistry32, Ownable {
         params.totalOperatorStake = getOperatorStake(operator, epoch);
 
         if (params.totalOperatorStake < amount) {
-            revert TooBigSlashAmount();
+            revert Middleware__TooBigSlashAmount();
         }
 
         // simple pro-rata slasher
-        for (uint256 i; i < vaults.length(); ++i) {
-            (address vault, uint48 enabledTime, uint48 disabledTime) = vaults.atWithTimes(i);
-            console2.log("Vault: ", vault);
-            console2.log("Enabled: ", enabledTime);
-            console2.log("Disabled: ", disabledTime);
-            console2.log("Epoch: ", params.epochStartTs);
-            console2.log("Active: ", _wasActiveAt(enabledTime, disabledTime, params.epochStartTs));
+        for (uint256 i; i < s_vaults.length(); ++i) {
+            (address vault, uint48 enabledTime, uint48 disabledTime) = s_vaults.atWithTimes(i);
             // just skip the vault if it was enabled after the target epoch or not enabled
             if (!_wasActiveAt(enabledTime, disabledTime, params.epochStartTs)) {
-                console2.log("skip vault", vault);
                 continue;
             }
 
@@ -346,9 +358,14 @@ contract Middleware is SimpleKeyRegistry32, Ownable {
         }
     }
 
+    /**
+     * @dev Get vault stake and calculate slashing amount.
+     * @param vault The vault address to calculate its stake
+     * @param params Struct containing slashing parameters
+     */
     function _processVaultSlashing(address vault, SlashParams memory params) private {
-        for (uint96 j = 0; j < subnetworksCnt; ++j) {
-            bytes32 subnetwork = NETWORK.subnetwork(j);
+        for (uint96 j = 0; j < s_subnetworksCount; ++j) {
+            bytes32 subnetwork = i_network.subnetwork(j);
             uint256 vaultStake = IBaseDelegator(IVault(vault).delegator()).stakeAt(
                 subnetwork, params.operator, params.epochStartTs, new bytes(0)
             );
@@ -358,22 +375,53 @@ contract Middleware is SimpleKeyRegistry32, Ownable {
         }
     }
 
-    function calcAndCacheStakes(
+    /**
+     * @dev Slashes a vault's stake for a specific operator
+     * @param timestamp Time at which the epoch started
+     * @param vault Address of the vault to slash
+     * @param subnetwork Subnetwork identifier
+     * @param operator Address of the operator being slashed
+     * @param amount Amount to slash
+     */
+    function _slashVault(
+        uint48 timestamp,
+        address vault,
+        bytes32 subnetwork,
+        address operator,
+        uint256 amount
+    ) private {
+        address slasher = IVault(vault).slasher();
+        uint256 slasherType = IEntity(slasher).TYPE();
+        if (slasherType == INSTANT_SLASHER_TYPE) {
+            ISlasher(slasher).slash(subnetwork, operator, amount, timestamp, new bytes(0));
+        } else if (slasherType == VETO_SLASHER_TYPE) {
+            IVetoSlasher(slasher).requestSlash(subnetwork, operator, amount, timestamp, new bytes(0));
+        } else {
+            revert Middleware__UnknownSlasherType();
+        }
+    }
+
+    /**
+     * @dev Calculates total stake for an epoch
+     * @param epoch The epoch to calculate stake for
+     * @return totalStake The total stake amount
+     */
+    function _calcTotalStake(
         uint48 epoch
-    ) public returns (uint256 totalStake) {
+    ) private view returns (uint256 totalStake) {
         uint48 epochStartTs = getEpochStartTs(epoch);
 
-        // for epoch older than SLASHING_WINDOW total stake can be invalidated (use cache)
-        if (epochStartTs < Time.timestamp() - SLASHING_WINDOW) {
-            revert TooOldEpoch();
+        // for epoch older than i_slashingWindow total stake can be invalidated (use cache)
+        if (epochStartTs < Time.timestamp() - i_slashingWindow) {
+            revert Middleware__TooOldEpoch();
         }
 
         if (epochStartTs > Time.timestamp()) {
-            revert InvalidEpoch();
+            revert Middleware__InvalidEpoch();
         }
 
-        for (uint256 i; i < operators.length(); ++i) {
-            (address operator, uint48 enabledTime, uint48 disabledTime) = operators.atWithTimes(i);
+        for (uint256 i; i < s_operators.length(); ++i) {
+            (address operator, uint48 enabledTime, uint48 disabledTime) = s_operators.atWithTimes(i);
 
             // just skip operator if it was added after the target epoch or paused
             if (!_wasActiveAt(enabledTime, disabledTime, epochStartTs)) {
@@ -381,61 +429,141 @@ contract Middleware is SimpleKeyRegistry32, Ownable {
             }
 
             uint256 operatorStake = getOperatorStake(operator, epoch);
-            operatorStakeCache[epoch][operator] = operatorStake;
-
             totalStake += operatorStake;
         }
+    }
 
-        totalStakeCached[epoch] = true;
-        totalStakeCache[epoch] = totalStake;
+    /**
+     * @dev Checks if an entity was active at a specific timestamp
+     * @param enabledTime Time when entity was enabled
+     * @param disabledTime Time when entity was disabled (0 if never disabled)
+     * @param timestamp Timestamp to check activity for
+     * @return bool True if entity was active at timestamp
+     */
+    function _wasActiveAt(uint48 enabledTime, uint48 disabledTime, uint48 timestamp) private pure returns (bool) {
+        return enabledTime != 0 && enabledTime <= timestamp && (disabledTime == 0 || disabledTime >= timestamp);
     }
 
-    function _calcTotalStake(
-        uint48 epoch
-    ) private view returns (uint256 totalStake) {
+    /**
+     * @notice Checks if a vault is registered
+     * @param vault The vault address to check
+     * @return bool True if vault is registered
+     */
+    function isVaultRegistered(
+        address vault
+    ) external view returns (bool) {
+        return s_vaults.contains(vault);
+    }
+
+    /**
+     * @notice Gets operator's stake for an epoch
+     * @param operator The operator address
+     * @param epoch The epoch number
+     * @return stake The operator's total stake
+     */
+    function getOperatorStake(address operator, uint48 epoch) public view returns (uint256 stake) {
+        if (s_totalStakeCached[epoch]) {
+            return s_operatorStakeCache[epoch][operator];
+        }
+
         uint48 epochStartTs = getEpochStartTs(epoch);
+        for (uint256 i; i < s_vaults.length(); ++i) {
+            (address vault, uint48 enabledTime, uint48 disabledTime) = s_vaults.atWithTimes(i);
 
-        // for epoch older than SLASHING_WINDOW total stake can be invalidated (use cache)
-        if (epochStartTs < Time.timestamp() - SLASHING_WINDOW) {
-            revert TooOldEpoch();
+            // just skip the vault if it was enabled after the target epoch or not enabled
+            if (!_wasActiveAt(enabledTime, disabledTime, epochStartTs)) {
+                continue;
+            }
+
+            for (uint96 j = 0; j < s_subnetworksCount; ++j) {
+                stake += IBaseDelegator(IVault(vault).delegator()).stakeAt(
+                    i_network.subnetwork(j), operator, epochStartTs, new bytes(0)
+                );
+            }
         }
 
-        if (epochStartTs > Time.timestamp()) {
-            revert InvalidEpoch();
+        return stake;
+    }
+
+    /**
+     * @notice Gets total stake for an epoch
+     * @param epoch The epoch number
+     * @return Total stake amount
+     */
+    function getTotalStake(
+        uint48 epoch
+    ) public view returns (uint256) {
+        if (s_totalStakeCached[epoch]) {
+            return s_totalStakeCache[epoch];
         }
+        return _calcTotalStake(epoch);
+    }
+
+    /**
+     * @notice Gets validator set for an epoch
+     * @param epoch The epoch number
+     * @return validatorsData Array of validator data
+     */
+    function getValidatorSet(
+        uint48 epoch
+    ) public view returns (ValidatorData[] memory validatorsData) {
+        uint48 epochStartTs = getEpochStartTs(epoch);
 
-        for (uint256 i; i < operators.length(); ++i) {
-            (address operator, uint48 enabledTime, uint48 disabledTime) = operators.atWithTimes(i);
+        validatorsData = new ValidatorData[](s_operators.length());
+        uint256 valIdx = 0;
+
+        for (uint256 i; i < s_operators.length(); ++i) {
+            (address operator, uint48 enabledTime, uint48 disabledTime) = s_operators.atWithTimes(i);
 
             // just skip operator if it was added after the target epoch or paused
             if (!_wasActiveAt(enabledTime, disabledTime, epochStartTs)) {
                 continue;
             }
 
-            uint256 operatorStake = getOperatorStake(operator, epoch);
-            totalStake += operatorStake;
+            bytes32 key = getOperatorKeyAt(operator, epochStartTs);
+            if (key == bytes32(0)) {
+                continue;
+            }
+
+            uint256 stake = getOperatorStake(operator, epoch);
+
+            validatorsData[valIdx++] = ValidatorData(stake, key);
+        }
+
+        // shrink array to skip unused slots
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(validatorsData, valIdx)
         }
     }
 
-    function _wasActiveAt(uint48 enabledTime, uint48 disabledTime, uint48 timestamp) private pure returns (bool) {
-        return enabledTime != 0 && enabledTime <= timestamp && (disabledTime == 0 || disabledTime >= timestamp);
+    /**
+     * @notice Gets the timestamp when an epoch starts
+     * @param epoch The epoch number
+     * @return timestamp The start time of the epoch
+     */
+    function getEpochStartTs(
+        uint48 epoch
+    ) public view returns (uint48 timestamp) {
+        return i_startTime + epoch * i_epochDuration;
     }
 
-    function _slashVault(
-        uint48 timestamp,
-        address vault,
-        bytes32 subnetwork,
-        address operator,
-        uint256 amount
-    ) private {
-        address slasher = IVault(vault).slasher();
-        uint256 slasherType = IEntity(slasher).TYPE();
-        if (slasherType == INSTANT_SLASHER_TYPE) {
-            ISlasher(slasher).slash(subnetwork, operator, amount, timestamp, new bytes(0));
-        } else if (slasherType == VETO_SLASHER_TYPE) {
-            IVetoSlasher(slasher).requestSlash(subnetwork, operator, amount, timestamp, new bytes(0));
-        } else {
-            revert UnknownSlasherType();
-        }
+    /**
+     * @notice Determines which epoch a timestamp belongs to
+     * @param timestamp The timestamp to check
+     * @return epoch The corresponding epoch number
+     */
+    function getEpochAtTs(
+        uint48 timestamp
+    ) public view returns (uint48 epoch) {
+        return (timestamp - i_startTime) / i_epochDuration;
+    }
+
+    /**
+     * @notice Gets the current epoch number
+     * @return epoch The current epoch
+     */
+    function getCurrentEpoch() public view returns (uint48 epoch) {
+        return getEpochAtTs(Time.timestamp());
     }
 }
diff --git a/test/MapWithTimeData.t.sol b/test/MapWithTimeData.t.sol
index ed66149..b848153 100644
--- a/test/MapWithTimeData.t.sol
+++ b/test/MapWithTimeData.t.sol
@@ -1,4 +1,17 @@
-// SPDX-License-Identifier: MIT
+//SPDX-License-Identifier: GPL-3.0-or-later
+
+// Copyright (C) Moondance Labs Ltd.
+// This file is part of Tanssi.
+// Tanssi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// Tanssi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License
+// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
 pragma solidity 0.8.25;
 
 import {Test, console2} from "forge-std/Test.sol";
diff --git a/test/mocks/MapWithTimeDataContract.sol b/test/mocks/MapWithTimeDataContract.sol
index eae5fd9..1323925 100644
--- a/test/mocks/MapWithTimeDataContract.sol
+++ b/test/mocks/MapWithTimeDataContract.sol
@@ -1,4 +1,17 @@
-// SPDX-License-Identifier: MIT
+//SPDX-License-Identifier: GPL-3.0-or-later
+
+// Copyright (C) Moondance Labs Ltd.
+// This file is part of Tanssi.
+// Tanssi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// Tanssi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License
+// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
 pragma solidity 0.8.25;
 
 import {EnumerableMap} from "@openzeppelin/contracts/utils/structs/EnumerableMap.sol";
diff --git a/test/mocks/symbiotic/DelegatorMock.sol b/test/mocks/symbiotic/DelegatorMock.sol
index ecf6d4a..cab0654 100644
--- a/test/mocks/symbiotic/DelegatorMock.sol
+++ b/test/mocks/symbiotic/DelegatorMock.sol
@@ -1,4 +1,17 @@
-// SPDX-License-Identifier: MIT
+//SPDX-License-Identifier: GPL-3.0-or-later
+
+// Copyright (C) Moondance Labs Ltd.
+// This file is part of Tanssi.
+// Tanssi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// Tanssi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License
+// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
 pragma solidity 0.8.25;
 
 import {console2} from "forge-std/console2.sol";
@@ -26,9 +39,9 @@ contract DelegatorMock is BaseDelegator {
     {}
 
     function _stakeAt(
-        bytes32 subnetwork,
+        bytes32, /*subnetwork*/
         address operator,
-        uint48 timestamp,
+        uint48, /*timestamp*/
         bytes memory hints
     ) internal view override returns (uint256, bytes memory) {
         uint256 operatorStake = IVault(vault).activeBalanceOf(operator);
diff --git a/test/mocks/symbiotic/OptInServiceMock.sol b/test/mocks/symbiotic/OptInServiceMock.sol
index 7c629e4..70baed2 100644
--- a/test/mocks/symbiotic/OptInServiceMock.sol
+++ b/test/mocks/symbiotic/OptInServiceMock.sol
@@ -1,4 +1,17 @@
-// SPDX-License-Identifier: MIT
+//SPDX-License-Identifier: GPL-3.0-or-later
+
+// Copyright (C) Moondance Labs Ltd.
+// This file is part of Tanssi.
+// Tanssi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// Tanssi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License
+// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
 pragma solidity 0.8.25;
 
 import {console2} from "forge-std/console2.sol";
@@ -46,7 +59,7 @@ contract OptInServiceMock is EIP712, IOptInService {
         _optIn(msg.sender, where);
     }
 
-    function optIn(address who, address where, uint48 deadline, bytes calldata signature) external {
+    function optIn(address who, address where, uint48, /*deadline*/ bytes calldata /*signature*/ ) external {
         _optIn(who, where);
     }
 
@@ -56,7 +69,7 @@ contract OptInServiceMock is EIP712, IOptInService {
         _optOut(msg.sender, where);
     }
 
-    function optOut(address who, address where, uint48 deadline, bytes calldata signature) external {
+    function optOut(address who, address where, uint48, /*deadline*/ bytes calldata /* signature*/ ) external {
         _optOut(who, where);
     }
 
diff --git a/test/mocks/symbiotic/RegistryMock.sol b/test/mocks/symbiotic/RegistryMock.sol
index c36725d..13eed4d 100644
--- a/test/mocks/symbiotic/RegistryMock.sol
+++ b/test/mocks/symbiotic/RegistryMock.sol
@@ -1,4 +1,17 @@
-// SPDX-License-Identifier: MIT
+//SPDX-License-Identifier: GPL-3.0-or-later
+
+// Copyright (C) Moondance Labs Ltd.
+// This file is part of Tanssi.
+// Tanssi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// Tanssi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License
+// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
 pragma solidity 0.8.25;
 
 import {Registry} from "@symbiotic/contracts/common/Registry.sol";
diff --git a/test/mocks/symbiotic/VaultMock.sol b/test/mocks/symbiotic/VaultMock.sol
index db5c3d9..93c8874 100644
--- a/test/mocks/symbiotic/VaultMock.sol
+++ b/test/mocks/symbiotic/VaultMock.sol
@@ -1,4 +1,17 @@
-// SPDX-License-Identifier: MIT
+//SPDX-License-Identifier: GPL-3.0-or-later
+
+// Copyright (C) Moondance Labs Ltd.
+// This file is part of Tanssi.
+// Tanssi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// Tanssi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License
+// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
 pragma solidity 0.8.25;
 
 import {console2} from "forge-std/console2.sol";
@@ -29,8 +42,8 @@ contract VaultMock is VaultStorage, MigratableEntity, IVault {
 
     function activeBalanceOfAt(
         address account,
-        uint48 timestamp,
-        bytes calldata hints
+        uint48, /*timestamp*/
+        bytes calldata /*hints*/
     ) external view returns (uint256) {
         return operatorStake[account];
     }
@@ -66,7 +79,7 @@ contract VaultMock is VaultStorage, MigratableEntity, IVault {
 
     function claimBatch(address recipient, uint256[] calldata epochs) external returns (uint256 amount) {}
 
-    function onSlash(uint256 amount, uint48 captureTimestamp) external returns (uint256 slashedAmount) {
+    function onSlash(uint256 amount, uint48 /*captureTimestamp*/ ) external returns (uint256 slashedAmount) {
         totalAtStake -= amount;
         slashedAmount = amount;
         for (uint256 i = 0; i < operators.length; i++) {
diff --git a/test/unit/Middleware.t.sol b/test/unit/Middleware.t.sol
index 1349c46..f94fc65 100644
--- a/test/unit/Middleware.t.sol
+++ b/test/unit/Middleware.t.sol
@@ -1,4 +1,17 @@
-// SPDX-License-Identifier: UNLICENSED
+//SPDX-License-Identifier: GPL-3.0-or-later
+
+// Copyright (C) Moondance Labs Ltd.
+// This file is part of Tanssi.
+// Tanssi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// Tanssi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// You should have received a copy of the GNU General Public License
+// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
 pragma solidity ^0.8.13;
 
 import {Test, console2} from "forge-std/Test.sol";
@@ -120,7 +133,7 @@ contract MiddlewareTest is Test {
         uint48 SHORT_SLASHING_WINDOW_ = 99;
 
         vm.startPrank(owner);
-        vm.expectRevert(Middleware.SlashingWindowTooShort.selector);
+        vm.expectRevert(Middleware.Middleware__SlashingWindowTooShort.selector);
 
         new Middleware(
             address(0),
@@ -197,12 +210,12 @@ contract MiddlewareTest is Test {
     }
 
     function testInitialState() public view {
-        assertEq(middleware.NETWORK(), address(network));
-        assertEq(middleware.OPERATOR_REGISTRY(), address(registry));
-        assertEq(middleware.VAULT_REGISTRY(), address(registry));
-        assertEq(middleware.EPOCH_DURATION(), NETWORK_EPOCH_DURATION);
-        assertEq(middleware.SLASHING_WINDOW(), SLASHING_WINDOW);
-        assertEq(middleware.subnetworksCnt(), 1);
+        assertEq(middleware.i_network(), address(network));
+        assertEq(middleware.i_operatorRegistry(), address(registry));
+        assertEq(middleware.i_vaultRegistry(), address(registry));
+        assertEq(middleware.i_epochDuration(), NETWORK_EPOCH_DURATION);
+        assertEq(middleware.i_slashingWindow(), SLASHING_WINDOW);
+        assertEq(middleware.s_subnetworksCount(), 1);
     }
 
     // ************************************************************************************************
@@ -234,7 +247,7 @@ contract MiddlewareTest is Test {
         vm.startPrank(owner);
         middleware.registerOperator(operator, OPERATOR_KEY);
 
-        vm.expectRevert(Middleware.OperatorAlreadyRegistred.selector);
+        vm.expectRevert(Middleware.Middleware__OperatorAlreadyRegistred.selector);
         middleware.registerOperator(operator, OPERATOR_KEY);
         vm.stopPrank();
     }
@@ -243,7 +256,7 @@ contract MiddlewareTest is Test {
         _registerOperatorToNetwork(operator, address(vault), true, false);
 
         vm.startPrank(owner);
-        vm.expectRevert(Middleware.NotOperator.selector);
+        vm.expectRevert(Middleware.Middleware__NotOperator.selector);
         middleware.registerOperator(owner, OPERATOR_KEY);
         vm.stopPrank();
     }
@@ -252,7 +265,7 @@ contract MiddlewareTest is Test {
         _registerOperatorToNetwork(operator, address(vault), false, true);
 
         vm.startPrank(owner);
-        vm.expectRevert(Middleware.OperatorNotOptedIn.selector);
+        vm.expectRevert(Middleware.Middleware__OperatorNotOptedIn.selector);
         middleware.registerOperator(operator, OPERATOR_KEY);
         vm.stopPrank();
     }
@@ -300,7 +313,7 @@ contract MiddlewareTest is Test {
 
     function testUpdateOperatorKeyNotRegistered() public {
         vm.startPrank(owner);
-        vm.expectRevert(Middleware.OperatorNotRegistred.selector);
+        vm.expectRevert(Middleware.Middleware__OperatorNotRegistred.selector);
         middleware.updateOperatorKey(operator, OPERATOR_KEY);
         vm.stopPrank();
     }
@@ -376,7 +389,7 @@ contract MiddlewareTest is Test {
 
         middleware.pauseOperator(operator);
         vm.warp(START_TIME + SLASHING_WINDOW - 1);
-        vm.expectRevert(Middleware.OperarorGracePeriodNotPassed.selector);
+        vm.expectRevert(Middleware.Middleware__OperarorGracePeriodNotPassed.selector);
         middleware.unregisterOperator(operator);
         vm.stopPrank();
     }
@@ -411,14 +424,14 @@ contract MiddlewareTest is Test {
 
         middleware.registerVault(address(vault));
 
-        vm.expectRevert(Middleware.VaultAlreadyRegistered.selector);
+        vm.expectRevert(Middleware.Middleware__VaultAlreadyRegistered.selector);
         middleware.registerVault(address(vault));
         vm.stopPrank();
     }
 
     function testRegisterVaultNotVault() public {
         vm.startPrank(owner);
-        vm.expectRevert(Middleware.NotVault.selector);
+        vm.expectRevert(Middleware.Middleware__NotVault.selector);
         middleware.registerVault(owner);
         vm.stopPrank();
     }
@@ -429,7 +442,7 @@ contract MiddlewareTest is Test {
         vm.startPrank(owner);
         vault.setSlasher(address(vetoSlasher));
         vm.store(address(vetoSlasher), bytes32(uint256(0)), bytes32(uint256(uint160(address(vault)))));
-        vm.expectRevert(Middleware.VaultEpochTooShort.selector);
+        vm.expectRevert(Middleware.Middleware__VaultEpochTooShort.selector);
         middleware.registerVault(address(vault));
         vm.stopPrank();
     }
@@ -524,7 +537,7 @@ contract MiddlewareTest is Test {
 
         middleware.pauseVault(address(vault));
         vm.warp(START_TIME + SLASHING_WINDOW - 1);
-        vm.expectRevert(Middleware.VaultGracePeriodNotPassed.selector);
+        vm.expectRevert(Middleware.Middleware__VaultGracePeriodNotPassed.selector);
         middleware.unregisterVault(address(vault));
         vm.stopPrank();
     }
@@ -535,31 +548,31 @@ contract MiddlewareTest is Test {
 
     function testSetSubnetworksCnt() public {
         vm.startPrank(owner);
-        middleware.setSubnetworksCnt(2);
-        assertEq(middleware.subnetworksCnt(), 2);
+        middleware.setSubnetworksCount(2);
+        assertEq(middleware.s_subnetworksCount(), 2);
         vm.stopPrank();
     }
 
     function testSetSubnetworksCntInvalidIfGreaterThanZero() public {
         vm.startPrank(owner);
-        middleware.setSubnetworksCnt(10);
-        assertEq(middleware.subnetworksCnt(), 10);
+        middleware.setSubnetworksCount(10);
+        assertEq(middleware.s_subnetworksCount(), 10);
 
-        vm.expectRevert(Middleware.InvalidSubnetworksCnt.selector);
-        middleware.setSubnetworksCnt(8);
+        vm.expectRevert(Middleware.Middleware__InvalidSubnetworksCnt.selector);
+        middleware.setSubnetworksCount(8);
         vm.stopPrank();
     }
 
     function testSetSubnetworksCntInvalid() public {
         vm.startPrank(owner);
-        vm.expectRevert(Middleware.InvalidSubnetworksCnt.selector);
-        middleware.setSubnetworksCnt(0);
+        vm.expectRevert(Middleware.Middleware__InvalidSubnetworksCnt.selector);
+        middleware.setSubnetworksCount(0);
         vm.stopPrank();
     }
 
     function testSetSubnetworksCntUnauthorized() public {
         vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, address(this)));
-        middleware.setSubnetworksCnt(2);
+        middleware.setSubnetworksCount(2);
     }
 
     // ************************************************************************************************
@@ -733,7 +746,7 @@ contract MiddlewareTest is Test {
         vm.warp(START_TIME + SLASHING_WINDOW + 1);
         uint48 currentEpoch = middleware.getCurrentEpoch();
         vm.warp(SLASHING_WINDOW * 2 + 1);
-        vm.expectRevert(Middleware.TooOldEpoch.selector);
+        vm.expectRevert(Middleware.Middleware__TooOldEpoch.selector);
         middleware.getTotalStake(currentEpoch);
         vm.stopPrank();
     }
@@ -750,7 +763,7 @@ contract MiddlewareTest is Test {
         vm.warp(START_TIME + SLASHING_WINDOW + 1);
         uint48 currentEpoch = middleware.getCurrentEpoch();
         vm.warp(START_TIME + SLASHING_WINDOW - 1);
-        vm.expectRevert(Middleware.InvalidEpoch.selector);
+        vm.expectRevert(Middleware.Middleware__InvalidEpoch.selector);
         middleware.getTotalStake(currentEpoch + 1);
         vm.stopPrank();
     }
@@ -874,7 +887,7 @@ contract MiddlewareTest is Test {
         vm.startPrank(owner);
         uint48 currentEpoch = middleware.getCurrentEpoch();
         vm.warp(SLASHING_WINDOW * 2 + 1);
-        vm.expectRevert(Middleware.TooOldEpoch.selector);
+        vm.expectRevert(Middleware.Middleware__TooOldEpoch.selector);
         middleware.slash(currentEpoch, operator, OPERATOR_STAKE);
         vm.stopPrank();
     }
@@ -893,7 +906,7 @@ contract MiddlewareTest is Test {
         uint256 totalStakeCached = middleware.calcAndCacheStakes(currentEpoch);
 
         uint256 slashAmount = OPERATOR_STAKE * 2;
-        vm.expectRevert(Middleware.TooBigSlashAmount.selector);
+        vm.expectRevert(Middleware.Middleware__TooBigSlashAmount.selector);
         middleware.slash(currentEpoch, operator, slashAmount);
 
         uint256 totalStake = middleware.getTotalStake(currentEpoch);
@@ -999,7 +1012,7 @@ contract MiddlewareTest is Test {
         uint48 currentEpoch = middleware.getCurrentEpoch();
 
         uint256 slashAmount = OPERATOR_STAKE / 2;
-        vm.expectRevert(Middleware.UnknownSlasherType.selector);
+        vm.expectRevert(Middleware.Middleware__UnknownSlasherType.selector);
         middleware.slash(currentEpoch, operator, slashAmount);
 
         vm.stopPrank();
@@ -1044,7 +1057,7 @@ contract MiddlewareTest is Test {
         vm.warp(START_TIME + SLASHING_WINDOW + 1);
         uint48 currentEpoch = middleware.getCurrentEpoch();
         vm.warp(SLASHING_WINDOW * 2 + 1);
-        vm.expectRevert(Middleware.TooOldEpoch.selector);
+        vm.expectRevert(Middleware.Middleware__TooOldEpoch.selector);
         middleware.calcAndCacheStakes(currentEpoch);
         vm.stopPrank();
     }
@@ -1062,7 +1075,7 @@ contract MiddlewareTest is Test {
         vm.warp(START_TIME + SLASHING_WINDOW + 1);
         uint48 currentEpoch = middleware.getCurrentEpoch();
         vm.warp(START_TIME + SLASHING_WINDOW - 1);
-        vm.expectRevert(Middleware.InvalidEpoch.selector);
+        vm.expectRevert(Middleware.Middleware__InvalidEpoch.selector);
         middleware.calcAndCacheStakes(currentEpoch + 1);
         vm.stopPrank();
     }