diff --git a/CHANGELOG.md b/CHANGELOG.md
index d1782103..00b6edec 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -48,6 +48,8 @@ This format is based on [Keep A Changelog](https://keepachangelog.com/en/1.0.0).
     - `initLovelaceAssertValue`
     - `initLovelaceAssertValueWith`
     - `withCollateral`
+  - Initialising wallets with staking keys.
+  - Access to initialised wallets via `WalletLookups`
 
 ## 0.1 -- 2022-02-14
 
diff --git a/README.md b/README.md
index eedac23a..7d9ed89a 100644
--- a/README.md
+++ b/README.md
@@ -40,6 +40,7 @@ NOTE: This branch launches local network in `Vasil`. It was tested with node `1.
 
 ## Tutorials
 
+* NEW: [Address types](docs/wallet-tags-and-addresses.md)
 * [Running disposable local network and building own runners](./local-cluster/README.md)
 * [Tasty integration](./docs/tasty-integration.md)
 * [Running Contracts is REPL](./docs/interactive-plutip.md)
diff --git a/docs/interactive-plutip.md b/docs/interactive-plutip.md
index 9b41a85b..18300be7 100644
--- a/docs/interactive-plutip.md
+++ b/docs/interactive-plutip.md
@@ -43,7 +43,10 @@ setup = do
   pure (env, ownWallet)
 
 addWalletWithAdas :: [Ada] -> ReaderT ClusterEnv IO BpiWallet
-addWalletWithAdas = addSomeWallet . map (fromInteger . Ada.toLovelace)
+addWalletWithAdas funds = 
+  addSomeWallet
+    (EntTag "w1")
+    (map (fromInteger . Ada.toLovelace) funds)
 ```
 
 > Aside: Feel free to choose the amount of ada you want to fund your wallet with. Just remember: `addSomeWallet` takes a list of _lovelace_ amounts. Here, I've actually made my custom `Ada` type as well some helper utilities (not the same as `Plutus.V1.Ledger.Ada` as that is removed in newer `plutus-ledger-api` versions).
@@ -179,7 +182,7 @@ import Test.Plutip.Contract.Types (TestContractConstraints)
 newtype ContractRunner = ContrRunner
   { runContr ::
       forall w e a.
-      TestContractConstraints w e a =>
+      TestContractConstraints w e Int a =>
       Contract w EmptySchema e a ->
       IO (Either (FailureReason e) a)
   }
@@ -192,7 +195,7 @@ begin = do
     setup = do
       env <- ask
       -- Gotta have all those utxos for the collaterals.
-      ownWallet <- addWalletWithAdas $ 300 : replicate 50 10
+      ownWallet <- addWalletWithAdas (300 : replicate 50 10) 
       -- Wait for faucet funds to be added.
       waitSeconds 2
       pure (env, ownWallet)
diff --git a/docs/tasty-integration.md b/docs/tasty-integration.md
index 3a22b781..e475e554 100644
--- a/docs/tasty-integration.md
+++ b/docs/tasty-integration.md
@@ -10,8 +10,8 @@ test =
   withConfiguredCluster def -- 1
     "Basic integration: launch, add wallet, tx from wallet to wallet" -- 2
     $ [ assertExecution "Contract 1" -- 3
-          (initAda [100,200] <> initLovelace 10_000_000) -- 3.1
-          (withContract $ \[wallet2pkh] -> someContract) -- 3.2
+          (initAda (EntTag "w0") [100,200] <> initLovelace (BaseTag "w1") [10_000_000]) -- 3.1
+          (withContract $ \ws -> someContract) -- 3.2
           [ shouldSucceed -- 3.3
           ]
       ]
@@ -20,8 +20,8 @@ test =
 1. Will start the local network with default config (more on configuring below)
 2. Description of test group that will be run on current instance of the network
 3. Test scenario that will be performed on the local network with it's description. Scenario includes:
-   1. (3.1.) Initialization of `wallets`. In this case two addresses will be funded: first will have 2 UTxOs with 100 and 200 Ada, second - single UTxO with 10 Ada.
-   2. (3.2) Execution of "`someContract :: Contract w s e a`". `PaymentPubKeyHash` of *first wallet* will be accessible in `someContract` as "own PaymentPubKeyHash". e.g with `ownFirstPaymentPubKeyHash`. `PaymentPubKeyHash` of *second* initiated wallet is brought into scope by `wallet2pkh` during pattern match on list (more on that below).
+   1. (3.1.) Initialization of `wallets`. In this case two addresses will be funded: first - enterprise address - will have 2 UTxOs with 100 and 200 Ada, second - base address - single UTxO with 10 Ada.
+   2. (3.2) Execution of "`someContract :: Contract w s e a`". `PaymentPubKeyHash` of *first wallet* will be accessible in `someContract` as "own PaymentPubKeyHash". e.g with `ownFirstPaymentPubKeyHash`. `PaymentPubKeyHash` of *second* initiated wallet is accessible through `ws :: WalletLookups`  (more on that below).
    3. (3.3) List of checks or `predicates` which will be performed for the result of `someContract` execution.
 
 It is possible to run several scenarios on single network instance - note that `withConfiguredCluster` accepts list of `assertExecution`'s.
@@ -33,25 +33,54 @@ It is possible to initialize arbitrary number of `wallets` in second argument of
 E.g. if `wallets` initialized like
 
 ```haskell
-(initAda [100] <> initAda [200] <> initAda [300])
+(initAda (EntTag "w0") [100] <> initAda (EntTag "w1") [200] <> initAda (BaseTag "w2") [300])
 ```
 
-we will get 3 funded addresses represented by 3 corresponding `wallets`: 
+we will get 3 funded addresses represented by 3 corresponding `wallets`:
 
-* `PaymentPubKeyHash` of wallet `initAda [100]` will be "own" `PaymentPubKeyHash` for contract executed it test case.
-* `PaymentPubKeyHash` of `wallets` `initAda [200]` and `initAda [300]` will be available via lambda argument. I.e.:
+* `PaymentPubKeyHash` of wallet `w0` will be "own" `PaymentPubKeyHash` for contract executed it test case.
+* `PaymentPubKeyHash` of `wallets` `w1` and `w2` will be available via lambda wallet lookups argument. I.e.:
 
+```haskell
+withContract $ \ws -> do
+  EntWallet pkh1 <- lookupWallet ws (EntTag "w1")
+  BaseWallet pkh2 spkh2 <- lookupWallet ws (BaseTag "w2")
+  someContract
+```
+
+Note that the lookup return type depends on a query tag.
+
+* `pkh1` is `PaymentPubKeyHash` of `wallet` `initAda (EntTag "w1") [200]`
+* `pkh2` is `PaymentPubKeyHash` of `wallet` `initAda (BaseTag "w2") [300]` and `spkh2` is its `StakePubKeyHash`
+
+`PaymentPubKeyHash` of `wallet` `initAda (EntTag "w0") [100]` is meant to be `pkh0` and not presented in the lookups.
+
+You can execute a contract with base address as contracts address:
+
+```haskell
+(initAda (BaseTag "w0") [100])
+```
+
+and witness in contract
 
 ```haskell
-withContract $ \[pkh1, pkh2] -> someContract
+withContract $ \_ -> do
+  ourAddr :| _ <- Contract.ownAddresses
+  case ourAddr of
+    Address (PubKeyCredential ourPkh) (Just (StakingHash (PubKeyCredential ourSpkh))) -> logInfo "This is the address we will get."
+    _ -> error "Nothing else matters"
 ```
 
-where
+Use `mustPayToPubKeyAddress` instead of `mustPayToPubKey` when your address has staking keys.
 
-* `pkh1` is `PaymentPubKeyHash` of `wallet` `initAda [200]`
-* `pkh2` is `PaymentPubKeyHash` of `wallet` `initAda [300]`
+You can also query for wallet address right away:
 
-`PaymentPubKeyHash` of `wallet` `initAda [100]` is meant to be `pkh0` and not presented in the list.
+```haskell
+withContract $ \ws -> do
+  addr1 <- lookupAddress ws "w1"
+  addr2 <- lookupAddress ws "w2"
+  someContract
+```
 
 ## Executing contracts
 
@@ -59,13 +88,13 @@ It is possible to run arbitrary number of contracts in 3d argument of `assertExe
 
 ```haskell
   assertExecution "Some description"
-          ( initAda [100])
+          ( initAda (EntTag "w1") [100])
           ( do
               void $
                 withContract $
-                  \pkhs -> contract1
-              withContractAs 1 $
-                \pkhs -> contract2
+                  \ws -> contract1
+              withContractAs "w1" $
+                \ws -> contract2
           )
           [shouldSucceed]
 ```
@@ -85,33 +114,33 @@ For example, consider the following scenario:
 
 ```haskell
   assertExecution "Some description"
-          ( initAda [100] -- walletA
-            <> initAda [200] -- walletB
-            <> initAda [300] -- walletC
+          ( initAda (EntTag "a") [100] -- walletA
+            <> initAda (EntTag "b") [200] -- walletB
+            <> initAda (EntTag "c") [300] -- walletC
           )
           ( do
               void $
-                withContractAs 1 $ -- running contract with walletB
-                  \[walletA_PKH, walletC_PKH] -> setupContract1
+                withContractAs "b" $ -- running contract with walletB
+                  \ws -> do
+                    wallA <- lookupWallet ws (EntTag "a")
+                    setupContract1
               void $
-                withContractAs 2 $  -- running contract with walletC
-                  \[walletA_PKH, walletB_PKH] -> setupContract2
-              withContract $
-                \pkhs -> theContract
+                withContractAs "c" $  -- running contract with walletC
+                  \ws -> do 
+                    wallB <- lookupWallet ws (EntTag "b")
+                    setupContract2
+              withContract $  -- uses first wallet, walletA
+                \ws -> theContract
           )
           [shouldSucceed]
 ```
 
-Under the hood, test runner builds list of wallets like this `[walletA, walletB, walletC]` and by calling `withContractAs` we can refer to an index (0 based) of specific wallet in this list. In that case, `PaymentPubKeyHash` of referenced `wallet` becomes "own" `PaymentPubKeyHash` of the contract, and `PaymentPubKeyHash`'es in argument of lambda  will be rearranged. E.g. in case of  `withContractAs 1`:
-
-* `PaymentPubKeyHash` of `walletB` will become own `PaymentPubKeyHash`
-* argument of lambada will contain `PaymentPubKeyHash`'es of `walletA` and `walletC`.
-
-Actually, `withContract` is just shortcut for `withContractAs 0`.
+`withContractAs` asks explicitly for the name of a wallet to be used as contract's.
+Instead `withContract` uses the first wallet, first in the order of how the initializations are written.
 
 ## Assertions
 
-To assert the result of contract execution user specifies list of checks or `predicates` as 4th argument of `assertExecution`. There are several `predicates` provided by the library that could be found in `Test.Plutip.Contract` module. Existing `predicates` allows to make assertions on Contracts state (`w`), error (`e`) and result (`a`) (consider type `Contract w s e a`). 
+To assert the result of contract execution user specifies list of checks or `predicates` as 4th argument of `assertExecution`. There are several `predicates` provided by the library that could be found in `Test.Plutip.Contract` module. Existing `predicates` allows to make assertions on Contracts state (`w`), error (`e`) and result (`a`) (consider type `Contract w s e a`).
 
 There are also predicates for making assertions on scripts execution budgets (e.g. `budgetsFitUnder` or `assertOverallBudget`). But be aware, that budget of script submitted to private network can differ from testnet or mainnet, at least because different amount of input UTxOs could be added during balancing, so this assertions are mostly useful for rough estimation and regression testing.
 
@@ -148,7 +177,7 @@ E.g. scenario like this
         assertExecutionWith
           [ShowBudgets]
           "Lock then spend contract"
-          (initAda (replicate 3 300))
+          (initAda (EntTag "w1") (replicate 3 300))
           (withContract $ const lockThenSpend)
           [ shouldSucceed
           ]
diff --git a/docs/wallet-tags-and-addresses.md b/docs/wallet-tags-and-addresses.md
new file mode 100644
index 00000000..6261c545
--- /dev/null
+++ b/docs/wallet-tags-and-addresses.md
@@ -0,0 +1,71 @@
+# Address types and `WalletTag`
+
+There are [several types of addresses](https://docs.cardano.org/learn/cardano-addresses) on Cardano network. Currently Plutip supports two types of addresses in both [interactive](./interactive-plutip.md) and [tasty](tasty-integration.md) modes:
+
+* ***Enterprise Address*** - carry no stake rights, backed only by payment keys
+* ***Base Address*** - directly specifies the staking key, backed by both payment and staking keys
+
+To pick which type of address to create for wallet, Plutip provides `WalletTag`. `WalletTag` has two constructors:
+
+* `BaseTag Text` - will create wallet with `Base Address`
+* `EntTag Text` - will create wallet with `Enterprise Address`
+
+Wallet tag also gives wallet it's name (the `Text` argument of constructor) and can be used to lookup desired wallet by name. There is special functionality for `WalletTag` available tasty framework, but lets see simple example with local cluster eDSL first.
+
+## `WalletTag` in interactive mode and local cluster
+
+Whenever you use eDSL function like `addSomeWallet` ([example 1](interactive-plutip.md), [example 2](../local-cluster/README.md)), from now on you can specify `WalletTag` as first argument to pick what type of address wallet will have. The `Text` argument of constructor will be added as textual tag to created wallet. E.g.:
+
+```haskell
+main :: IO ()
+main = do
+  (st, _) <- startCluster def $ do
+    wallet1 <- addSomeWallet (BaseTag "wallet1") [100_000_000]
+    waitSeconds 2
+  stopCluster st
+```
+
+`wallet1` has type `BpiWallet`, and textual tag can be obtained with
+
+```haskell
+wTag :: Text
+wTag = bwTag wallet1
+```
+
+It could useful if you want to print some info about wallet and distinguish output by tag, or when you initialize several wallets with something like `mapM` - can use tags to enable some lookups:
+
+```haskell
+  let mkTag idx = EntTag $ T.pack $ "wallet" <> show idx
+  wallets <- 
+    for [0..42] $ \idx ->
+      addSomeWalletDir (mkTag idx) [100_000_000]
+  ...
+  let maybeMyWallet = find ((== "wallet13") . bwTag) wallets
+```
+
+## `WalletTag` and tasty Plutip
+
+In tasty integration it is also possible to pick address type now with functions like `initAda`. And on top of that new lookup system was added, that lets you find initialized wallets easier (no more pattern matching on lists and figuring out correct indexes):
+
+```haskell
+assertExecution
+  "Some test"
+  ( initAndAda (EntTag "w0") [100]
+    <> initAnd (EntTag "w1") [100]
+  )
+  ( do
+  void $ withContract $ \ws -> do
+    -- lookup wallet by it's tag
+    EntWallet pkh0 <- lookupWallet ws (EntTag "w0")
+    payToPubKey pkh0 10_000_000
+  -- select wallet with tag "w1" to be "own" wallet
+  withContractAs "w1" $ \ws -> do
+    -- lookup address by wallet tag
+    addr1 <- lookupAddress ws "w1"
+    payToPubKeyAddress addr1 10_000_000
+  )
+  [shouldSucceed]
+
+```
+
+For more examples see [tasty integration tutorial](./tasty-integration.md).
diff --git a/flake.lock b/flake.lock
index e92c99d8..7203db86 100644
--- a/flake.lock
+++ b/flake.lock
@@ -71,17 +71,17 @@
         "typed-protocols": "typed-protocols"
       },
       "locked": {
-        "lastModified": 1665413135,
-        "narHash": "sha256-ZUfqxC7+5pLtmD/e1ABsHKGu2+hi/JvWRORvfnf13y4=",
+        "lastModified": 1665991587,
+        "narHash": "sha256-YZVYzIHyQS2LC9ahrzbNwY6V/cGnNyKvtscKyYM/fGY=",
         "owner": "mlabs-haskell",
         "repo": "bot-plutus-interface",
-        "rev": "d6cf1e3686bc31bb2571c6feefbe28e3a2c8bb06",
+        "rev": "74e9d951a96eada55efa3ed0178f3f2a4cffd492",
         "type": "github"
       },
       "original": {
         "owner": "mlabs-haskell",
-        "ref": "d6cf1e3686bc31bb2571c6feefbe28e3a2c8bb06",
         "repo": "bot-plutus-interface",
+        "rev": "74e9d951a96eada55efa3ed0178f3f2a4cffd492",
         "type": "github"
       }
     },
diff --git a/flake.nix b/flake.nix
index 2fbcfcc9..e069f1be 100644
--- a/flake.nix
+++ b/flake.nix
@@ -11,7 +11,7 @@
       flake = false;
     };
     bot-plutus-interface.url =
-      "github:mlabs-haskell/bot-plutus-interface?ref=d6cf1e3686bc31bb2571c6feefbe28e3a2c8bb06";
+      "github:mlabs-haskell/bot-plutus-interface/74e9d951a96eada55efa3ed0178f3f2a4cffd492";
   };
 
   outputs =
diff --git a/hie.yaml b/hie.yaml
index e776c92f..0adc4d8d 100644
--- a/hie.yaml
+++ b/hie.yaml
@@ -6,3 +6,5 @@ cradle:
       component: "test-suite:plutip-tests"
     - path: "./local-cluster/"
       component: "exe:local-cluster"
+    - path: "./plutip-server/"
+      component: "exe:plutip-server"
diff --git a/local-cluster/Main.hs b/local-cluster/Main.hs
index 8d13d47c..ec641b45 100644
--- a/local-cluster/Main.hs
+++ b/local-cluster/Main.hs
@@ -6,10 +6,11 @@
 module Main (main) where
 
 import Control.Applicative (optional, (<**>))
-import Control.Monad (forM_, replicateM, void)
+import Control.Monad (forM_, void)
 import Control.Monad.IO.Class (liftIO)
 import Control.Monad.Reader (ReaderT (ReaderT))
 import Data.Default (def)
+import Data.Traversable (for)
 import Numeric.Positive (Positive)
 import Options.Applicative (Parser, helper, info)
 import Options.Applicative qualified as Options
@@ -17,15 +18,18 @@ import Test.Plutip.Config
   ( PlutipConfig (clusterWorkingDir),
     WorkingDirectory (Fixed, Temporary),
   )
-import Test.Plutip.Internal.BotPlutusInterface.Wallet (addSomeWalletDir, walletPkh)
+import Test.Plutip.Internal.BotPlutusInterface.Wallet (addSomeWalletDir)
 import Test.Plutip.Internal.Types (nodeSocket)
 import Test.Plutip.LocalCluster
   ( mkMainnetAddress,
     startCluster,
     stopCluster,
     waitSeconds,
+    walletPaymentPkh
   )
 import GHC.Natural (Natural)
+import Test.Plutip.Internal.BotPlutusInterface.Types (WalletTag(EntTag), BpiWallet (bwTag))
+import Data.Text qualified as T
 
 main :: IO ()
 main = do
@@ -43,7 +47,7 @@ main = do
         waitSeconds 2 -- let wallet Tx finish, it can take more time with bigger slot length
 
         separate
-        liftIO $ forM_ (zip ws [(1 :: Int) ..]) printWallet
+        liftIO $ forM_ ws printWallet
         printNodeRelatedInfo
         separate
 
@@ -65,12 +69,14 @@ main = do
         amt -> Right $ fromInteger . toInteger $ amt
 
     initWallets numWallets numUtxos amt dirWallets = do
-      replicateM (max 0 numWallets) $
-        addSomeWalletDir (replicate numUtxos amt) dirWallets
-
-    printWallet (w, n) = do
-      putStrLn $ "Wallet " ++ show n ++ " PKH: " ++ show (walletPkh w)
-      putStrLn $ "Wallet " ++ show n ++ " mainnet address: " ++ show (mkMainnetAddress w)
+      for [0..(max 0 numWallets - 1)] $ \idx ->
+        addSomeWalletDir (EntTag $ T.pack $ show idx) 
+          (replicate numUtxos amt)  dirWallets
+
+    printWallet w = do
+      putStrLn $ "Wallet " ++ show (bwTag w) ++ " PKH: " ++ show (walletPaymentPkh w)
+      putStrLn $ "Wallet " ++ show (bwTag w) ++ " mainnet address: " 
+                  ++ show (mkMainnetAddress w)
 
     toAda = (* 1_000_000)
 
diff --git a/local-cluster/README.md b/local-cluster/README.md
index 9599b154..08629935 100644
--- a/local-cluster/README.md
+++ b/local-cluster/README.md
@@ -52,7 +52,7 @@ main = do
      ask >>= \cEnv -> runContract cEnv wallet contract
      
   (st, _) <- startCluster def $ do
-    w <- addSomeWallet [100_000_000]
+    w <- addSomeWallet (BaseTag "wallet1") [100_000_000]
     waitSeconds 2
     result <- executeContract w someContract
     doSomething result
diff --git a/plutip-server/Api.hs b/plutip-server/Api.hs
index 2fc12a2d..61ae452f 100644
--- a/plutip-server/Api.hs
+++ b/plutip-server/Api.hs
@@ -65,7 +65,7 @@ server serverOptions =
     :<|> stopClusterHandler
 
 appServer :: Env -> Server Api
-appServer env@(Env {options}) =
+appServer env@Env {options} =
   hoistServer api appHandler (server options)
   where
     appHandler :: forall (a :: Type). AppM a -> Handler a
diff --git a/plutip-server/Api/Handlers.hs b/plutip-server/Api/Handlers.hs
index 7e2614ef..ce41916c 100644
--- a/plutip-server/Api/Handlers.hs
+++ b/plutip-server/Api/Handlers.hs
@@ -20,12 +20,15 @@ import Data.Traversable (for)
 import System.Directory (doesFileExist)
 import System.FilePath (replaceFileName)
 import Test.Plutip.Config (chainIndexPort, relayNodeLogs)
+import Test.Plutip.Internal.BotPlutusInterface.Keys (KeyPair (sKey))
 import Test.Plutip.Internal.BotPlutusInterface.Setup (keysDir)
-import Test.Plutip.Internal.BotPlutusInterface.Wallet (BpiWallet (signKey), addSomeWallet)
+import Test.Plutip.Internal.BotPlutusInterface.Types (WalletTag (BaseTag, EntTag))
+import Test.Plutip.Internal.BotPlutusInterface.Wallet (BpiWallet (payKeys), addSomeWallet)
 import Test.Plutip.Internal.LocalCluster (startCluster, stopCluster)
 import Test.Plutip.Internal.Types (ClusterEnv (runningNode))
 import Test.Plutip.LocalCluster (waitSeconds)
 import Types (
+  AddressType (Base, Enterprise),
   AppM,
   ClusterStartupFailureReason (
     ClusterIsRunningAlready,
@@ -40,6 +43,7 @@ import Types (
     privateKeys
   ),
   Env (status),
+  Key (addressType, funds),
   Lovelace (unLovelace),
   PrivateKey,
   ServerOptions (ServerOptions, nodeLogs),
@@ -52,13 +56,15 @@ import Types (
   StopClusterResponse (StopClusterFailure, StopClusterSuccess),
  )
 
+import Data.Text qualified as T
+
 startClusterHandler :: ServerOptions -> StartClusterRequest -> AppM StartClusterResponse
 startClusterHandler
   ServerOptions {nodeLogs}
   StartClusterRequest {keysToGenerate} = interpret $ do
     -- Check that lovelace amounts are positive
-    for_ keysToGenerate $ \lovelaceAmounts -> do
-      for_ lovelaceAmounts $ \lovelaces -> do
+    for_ keysToGenerate $ \key -> do
+      for_ (funds key) $ \lovelaces -> do
         unless (unLovelace lovelaces > 0) $ do
           throwError NegativeLovelaces
     statusMVar <- asks status
@@ -82,9 +88,9 @@ startClusterHandler
       setup :: ReaderT ClusterEnv IO (ClusterEnv, [BpiWallet])
       setup = do
         env <- ask
+        let tags = T.pack . show <$> [0 ..]
         wallets <- do
-          for keysToGenerate $ \lovelaceAmounts -> do
-            addSomeWallet (fromInteger . unLovelace <$> lovelaceAmounts)
+          for (zip tags keysToGenerate) $ \(idx, key) -> addWallet key idx
         waitSeconds 2 -- wait for transactions to submit
         pure (env, wallets)
       getNodeSocketFile (runningNode -> RunningNode conn _ _ _) = nodeSocketFile conn
@@ -92,8 +98,13 @@ startClusterHandler
         -- assumption is that node.config lies in the same directory as node.socket
         flip replaceFileName "node.config" . getNodeSocketFile
       getWalletPrivateKey :: BpiWallet -> PrivateKey
-      getWalletPrivateKey = Text.decodeUtf8 . Base16.encode . serialiseToCBOR . signKey
+      getWalletPrivateKey = Text.decodeUtf8 . Base16.encode . serialiseToCBOR . sKey . payKeys
       interpret = fmap (either ClusterStartupFailure id) . runExceptT
+      addWallet key tag =
+        let funds' = (fromInteger . unLovelace <$> funds key)
+         in case addressType key of
+              Base -> addSomeWallet (BaseTag tag) funds'
+              Enterprise -> addSomeWallet (EntTag tag) funds'
 
 stopClusterHandler :: StopClusterRequest -> AppM StopClusterResponse
 stopClusterHandler StopClusterRequest = do
@@ -104,4 +115,4 @@ stopClusterHandler StopClusterRequest = do
     else do
       statusTVar <- liftIO $ takeMVar statusMVar
       liftIO $ stopCluster statusTVar
-      pure $ StopClusterSuccess
+      pure StopClusterSuccess
diff --git a/plutip-server/Main.hs b/plutip-server/Main.hs
index d093aa2e..be9a83e8 100644
--- a/plutip-server/Main.hs
+++ b/plutip-server/Main.hs
@@ -1,5 +1,3 @@
-{-# LANGUAGE NumericUnderscores #-}
-
 module Main (main) where
 
 import Api (app)
diff --git a/plutip-server/Types.hs b/plutip-server/Types.hs
index 1458491c..899e6b87 100644
--- a/plutip-server/Types.hs
+++ b/plutip-server/Types.hs
@@ -8,6 +8,8 @@ module Types (
   Env (Env, status, options),
   ErrorMessage,
   Lovelace (unLovelace),
+  Key (addressType, funds),
+  AddressType (Base, Enterprise),
   PlutipServerError (PlutipServerError),
   PrivateKey,
   ServerOptions (ServerOptions, nodeLogs, port),
@@ -37,14 +39,7 @@ import Data.Text (Text)
 import GHC.Generics (Generic)
 import Network.Wai.Handler.Warp (Port)
 import Test.Plutip.Internal.BotPlutusInterface.Wallet (BpiWallet)
-import Test.Plutip.Internal.LocalCluster (
-  ClusterStatus (
-    ClusterClosed,
-    ClusterClosing,
-    ClusterStarted,
-    ClusterStarting
-  ),
- )
+import Test.Plutip.Internal.LocalCluster (ClusterStatus)
 import Test.Plutip.Internal.Types (ClusterEnv)
 import UnliftIO.STM (TVar)
 
@@ -86,6 +81,19 @@ instance Exception PlutipServerError
 
 type ErrorMessage = Text
 
+data Key = Key
+  { addressType :: AddressType
+  , funds :: [Lovelace]
+  }
+  deriving stock (Show, Eq, Generic)
+  deriving anyclass (ToJSON, FromJSON)
+
+data AddressType
+  = Base
+  | Enterprise
+  deriving stock (Show, Eq, Generic)
+  deriving anyclass (ToJSON, FromJSON)
+
 newtype Lovelace = Lovelace {unLovelace :: Integer}
   deriving stock (Show, Eq, Generic)
   deriving newtype (ToJSON, Num)
@@ -99,7 +107,7 @@ instance FromJSON Lovelace where
 
 newtype StartClusterRequest = StartClusterRequest
   { -- | Lovelace amounts for each UTXO of each wallet
-    keysToGenerate :: [[Lovelace]]
+    keysToGenerate :: [Key]
   }
   deriving stock (Show, Eq, Generic)
   deriving anyclass (FromJSON, ToJSON)
diff --git a/plutip.cabal b/plutip.cabal
index ddec6548..c18753b5 100644
--- a/plutip.cabal
+++ b/plutip.cabal
@@ -29,6 +29,7 @@ common common-imports
     , cardano-addresses
     , cardano-api
     , cardano-crypto
+    , cardano-crypto-class
     , cardano-crypto-wrapper
     , cardano-ledger-core
     , cardano-slotting
@@ -47,6 +48,7 @@ common common-imports
     , http-client
     , http-types
     , iohk-monitoring
+    , lens
     , memory
     , mtl
     , openapi3
@@ -145,6 +147,7 @@ library
     Test.Plutip.Contract.Types
     Test.Plutip.Contract.Values
     Test.Plutip.Internal.BotPlutusInterface.Keys
+    Test.Plutip.Internal.BotPlutusInterface.Lookups
     Test.Plutip.Internal.BotPlutusInterface.Run
     Test.Plutip.Internal.BotPlutusInterface.Setup
     Test.Plutip.Internal.BotPlutusInterface.Types
@@ -174,7 +177,6 @@ test-suite plutip-tests
   ghc-options:        -Wall -threaded -rtsopts
   build-depends:
     , base
-    , lens
     , plutip
     , tasty
     , tasty-expected-failure
@@ -247,5 +249,6 @@ executable local-cluster
     , optparse-applicative
     , plutip
     , positive
+    , text
 
   ghc-options:   -Wall -threaded -rtsopts
diff --git a/src/Test/Plutip/Contract.hs b/src/Test/Plutip/Contract.hs
index c985e0fb..c274adbd 100644
--- a/src/Test/Plutip/Contract.hs
+++ b/src/Test/Plutip/Contract.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE GADTs #-}
 {-# LANGUAGE LambdaCase #-}
 
 -- |
@@ -20,6 +21,8 @@
 --  `initLovelace`. In addition, the value in these wallets can be asserted after the contract
 --  execution with `initAdaAssertValue` or `initAndAssertAda`. When `initAdaAssertValue` or `initAndAssertAda` used
 --  to initiate wallets corresponding test case will be added automatically.
+--  A wallet is named with a tag, the tag contstructor specifies what type of wallet is to be initialized.
+--  Don't use the same name `k` for two wallets even with different tag constructors.
 --
 --  Each assertion in assertions list will become separate test case in `TestTree`,
 --  however Contract will be executed only once.
@@ -28,7 +31,7 @@
 --
 --    > assertExecution
 --    >   "Some Contract"                   -- Contract description
---    >   (initAda 100)                     -- wallets and initial funds for them (single wallet in this case)
+--    >   (initAda (EntTag "w1") 100)                     -- wallets and initial funds for them (single wallet in this case)
 --    >   (withContract $ \_ -> myContract) -- contract execution
 --    >   [ shouldSucceed                   -- list of assertions
 --    >   , not $ shouldYield someResult
@@ -37,7 +40,7 @@
 --
 --  To use multiple wallets, you can use the `Semigroup` instance of `TestWallets`. To reference the
 --  wallet inside the contract, the following callback function is used together with `withContract`:
---  @[PaymentPubKeyHash] -> Contract w s e a@.
+--  @WalletLookups k -> Contract w s e a@.
 --
 -- To display information useful for debugging together with test results use `assertExecutionWith`
 -- and provide it with options:
@@ -46,15 +49,16 @@
 --    - ShowTrace, for displaying contract execution trace
 --    - ShowTraceButOnlyContext, like ShowTrace but filter what to show
 --
---  Note that @[PaymentPubKeyHash]@ does not include the contract's own wallet,
+--  Note that @WalletLookups@ don't include the contract's own wallet,
 --  for that you can use `Plutus.Contract.ownPaymentPubKeyHash` inside the Contract monad.
 --
 --  When contract supplied to test with `withContract`,
 --  the 1st initiated wallet will be used as "own" wallet, e.g.:
 --
 --    > assertExecution  "Send some Ada"
---    >   (initAda 100 <> initAda 101 <> initAda 102)
---    >   (withContract $ \[pkh1, pkh2] ->
+--    >   (initAda (EntTag "w0") [100] <> initAda (EntTag "w1") [101] <> initAda (EntTag "w2") [102])
+--    >   (withContract $ \ws -> do
+--    >     EntWallet pkh1 <- lookupWallet ws (EntTag "w1")
 --    >     payToPubKey pkh1 (Ada.lovelaceValueOf amt))
 --    >   [shouldSucceed]
 --
@@ -63,18 +67,19 @@
 --  - 3 wallets will be initialised with 100, 101 and 102 Ada respectively
 --  - wallet with 100 Ada will be used as own wallet to run the contract
 --  - `pkh1` - `PaymentPubKeyHash` of wallet with 101 Ada
---  - `pkh2` - `PaymentPubKeyHash` of wallet with 102 Ada
 --
 --
---  When contract supplied to test with `withContractAs`, wallet with provided index (0 based)
+--  When contract supplied to test with `withContractAs`, wallet with provided name
 --  will be used as "own" wallet, e.g.:
 --
 --    > assertExecutionWith
 --    >   [ShowBudgets, ShowTraceButOnlyContext ContractLog Error]
 --    >   "Send some Ada"
---    >   (initAda 100 <> initAda 101 <> initAda 102)
---    >   (withContractAs 1 $ \[pkh0, pkh2] ->
---    >     payToPubKey pkh1 (Ada.lovelaceValueOf amt))
+--    >   (initAda (EntTag "pkh0") 100 <> initAda (EntTag "myOwnWallet") 101 <> initAda (EntTag "pkh2") 102)
+--    >   (withContractAs "myOwnWallet" $ \ws -> do
+--    >     EntWallet pkh0 <- lookupWallet ws (EntTag "pkh0")
+--    >     EntWallet pkh2 <- lookupWallet ws (EntTag "pkh2")
+--    >     payToPubKey pkh2 (Ada.lovelaceValueOf amt))
 --    >   [shouldSucceed]
 --
 --  Here:
@@ -91,13 +96,15 @@
 --
 --    > assertExecution
 --    >   "Two contracts one after another"
---    >   (initAda 100 <> initAda 101)
+--    >   (initAda (EntTag "w0") [100] <> initAda (EntTag "w1") [101])
 --    >   ( do
 --    >       void $ -- run something prior to the contract which result will be checked
---    >         withContract $
---    >           \[pkh1] -> payTo pkh1 10_000_000
---    >       withContractAs 1 $ -- run the contract which result will be checked
---    >         \[pkh1] -> payTo pkh1 10_000_000
+--    >         withContract $ \ws ->
+--    >           EntWallet pkh1 <- lookupWallet ws (EntTag "w1")
+--    >           payTo pkh1 10_000_000
+--    >       withContractAs "w1" $ \ws -> do  -- run the contract which result will be checked
+--    >         EntWallet pkh0 <- lookupWallet ws (EntTag "w0")
+--    >         payTo pkh0 10_000_000
 --    >   )
 --    >   [shouldSucceed]
 --
@@ -107,8 +114,6 @@ module Test.Plutip.Contract (
   withContract,
   withContractAs,
   -- Wallet initialisation
-  TestWallets (TestWallets, unTestWallets),
-  TestWallet (twInitDistribuition, twExpected),
   initAda,
   withCollateral,
   initAndAssertAda,
@@ -121,12 +126,13 @@ module Test.Plutip.Contract (
   initLovelaceAssertValue,
   initLovelaceAssertValueWith,
   -- Helpers
-  ledgerPaymentPkh,
-  ValueOrdering (VEq, VGt, VLt, VGEq, VLEq),
+  walletPaymentPkh,
   assertValues,
   assertExecution,
   assertExecutionWith,
   ada,
+  TestWallets,
+  ClusterTest (ClusterTest),
   -- Contract runners
   runContract,
   runContractWithLogLvl,
@@ -144,17 +150,19 @@ import Control.Arrow (left)
 import Control.Monad.Reader (MonadIO (liftIO), MonadReader (ask), ReaderT, runReaderT, void)
 import Data.Bool (bool)
 import Data.Kind (Type)
-import Data.List.NonEmpty (NonEmpty)
+import Data.List.NonEmpty (NonEmpty, toList)
 import Data.List.NonEmpty qualified as NonEmpty
+import Data.Map (Map)
+import Data.Map.Strict qualified as Map
 import Data.Maybe (isJust)
 import Data.Row (Row)
 import Data.Tagged (Tagged (Tagged))
+import Data.Text (Text)
 import Data.Text qualified as Text
-import Ledger (PaymentPubKeyHash)
-import Ledger.Address (pubKeyHashAddress)
 import Ledger.Value (Value)
 import Plutus.Contract (Contract, waitNSlots)
 import PlutusPrelude (render)
+import PlutusTx.These (These (That, These, This))
 import Prettyprinter (Doc, Pretty (pretty), vcat, (<+>))
 import Test.Plutip.Contract.Init (
   initAda,
@@ -172,13 +180,24 @@ import Test.Plutip.Contract.Init (
 import Test.Plutip.Contract.Types (
   TestContract (TestContract),
   TestContractConstraints,
-  TestWallet (twExpected, twInitDistribuition),
-  TestWallets (TestWallets, unTestWallets),
-  ValueOrdering (VEq, VGEq, VGt, VLEq, VLt),
  )
 import Test.Plutip.Contract.Values (assertValues, valueAt)
-import Test.Plutip.Internal.BotPlutusInterface.Run (runContract, runContractWithLogLvl)
-import Test.Plutip.Internal.BotPlutusInterface.Wallet (BpiWallet, ledgerPaymentPkh)
+import Test.Plutip.Internal.BotPlutusInterface.Lookups (WalletLookups, lookupsMap, makeWalletInfo, makeWalletLookups)
+import Test.Plutip.Internal.BotPlutusInterface.Run (
+  runContract,
+  runContractWithLogLvl,
+ )
+import Test.Plutip.Internal.BotPlutusInterface.Types (
+  BpiWallet (bwTag),
+  TestWallets,
+  WalletInfo,
+  getTag,
+  ownAddress,
+  twExpected,
+ )
+import Test.Plutip.Internal.BotPlutusInterface.Wallet (
+  walletPaymentPkh,
+ )
 import Test.Plutip.Internal.Types (
   ClusterEnv,
   ExecutionResult (contractLogs, outcome),
@@ -193,7 +212,11 @@ import Test.Tasty.HUnit (assertFailure, testCase)
 import Test.Tasty.Providers (IsTest (run, testOptions), TestTree, singleTest, testPassed)
 
 type TestRunner (w :: Type) (e :: Type) (a :: Type) =
-  ReaderT (ClusterEnv, NonEmpty BpiWallet) IO (ExecutionResult w e (a, NonEmpty Value))
+  ReaderT (ClusterEnv, NonEmpty BpiWallet) IO (ExecutionResult w e (a, Map Text Value))
+
+-- | A type for the output of `assertExecution`.
+-- `k` is existentially quantified to allow different key types in every test case.
+newtype ClusterTest = ClusterTest (TestWallets, IO (ClusterEnv, NonEmpty BpiWallet) -> TestTree)
 
 -- | When used with `withCluster`, builds `TestTree` from initial wallets distribution,
 --  Contract and list of assertions (predicates). Each assertion will be run as separate test case,
@@ -201,7 +224,7 @@ type TestRunner (w :: Type) (e :: Type) (a :: Type) =
 --
 -- > assertExecution
 -- >   "Some Contract"                   -- Contract description
--- >   (initAda 100)                     -- wallets and initial funds for them (single wallet in this case)
+-- >   (initAda (EntTag "w0") [100])                     -- wallets and initial funds for them (single wallet in this case)
 -- >   (withContract $ \_ -> myContract) -- contract execution
 -- >   [ shouldSucceed                   -- list of assertions
 -- >   , not $ shouldYield someResult
@@ -216,7 +239,7 @@ assertExecution ::
   TestWallets ->
   TestRunner w e a ->
   [Predicate w e a] ->
-  (TestWallets, IO (ClusterEnv, NonEmpty BpiWallet) -> TestTree)
+  ClusterTest
 assertExecution = assertExecutionWith mempty
 
 -- | Version of assertExecution parametrised with a list of extra TraceOption's.
@@ -232,9 +255,9 @@ assertExecutionWith ::
   TestWallets ->
   TestRunner w e a ->
   [Predicate w e a] ->
-  (TestWallets, IO (ClusterEnv, NonEmpty BpiWallet) -> TestTree)
+  ClusterTest
 assertExecutionWith options tag testWallets testRunner predicates =
-  (testWallets, toTestGroup)
+  ClusterTest (testWallets, toTestGroup)
   where
     toTestGroup :: IO (ClusterEnv, NonEmpty BpiWallet) -> TestTree
     toTestGroup ioEnv =
@@ -247,11 +270,11 @@ assertExecutionWith options tag testWallets testRunner predicates =
               ((toCase ioRes <$> predicates) <> ((`optionToTestTree` ioRes) <$> options))
 
     -- wraps IO with result of contract execution into single test
-    toCase :: IO (ExecutionResult w e (a, NonEmpty Value)) -> Predicate w e a -> TestTree
+    toCase :: IO (ExecutionResult w e (a, Map Text Value)) -> Predicate w e a -> TestTree
     toCase ioRes p =
       singleTest (pTag p) (TestContract p ioRes)
 
-    optionToTestTree :: TraceOption -> IO (ExecutionResult w e (a, NonEmpty Value)) -> TestTree
+    optionToTestTree :: TraceOption -> IO (ExecutionResult w e (a, Map Text Value)) -> TestTree
     optionToTestTree = \case
       ShowBudgets -> singleTest "Budget stats" . StatsReport
       ShowTrace -> singleTest logsName . LogsReport DisplayAllTrace
@@ -265,15 +288,15 @@ assertExecutionWith options tag testWallets testRunner predicates =
 --
 -- @since 0.2
 maybeAddValuesCheck ::
-  Show e =>
-  IO (ExecutionResult w e (a, NonEmpty Value)) ->
+  (Show e) =>
+  IO (ExecutionResult w e (a, Map Text Value)) ->
   TestWallets ->
   [TestTree] ->
   [TestTree]
 maybeAddValuesCheck ioRes tws =
   bool id (valuesCheckCase :) (any isJust expected)
   where
-    expected = twExpected <$> unTestWallets tws
+    expected = Map.fromList $ toList $ (\tw -> (getTag tw, twExpected tw)) <$> tws
 
     valuesCheckCase :: TestTree
     valuesCheckCase =
@@ -288,57 +311,66 @@ maybeAddValuesCheck ioRes tws =
 
     checkValues o =
       left (Text.pack . show) o
-        >>= \(_, vs) -> assertValues expected vs
+        >>= \(_, vs) ->
+          let theseToPair = \case
+                (These b c) -> (b, c)
+                _ -> error "The two maps should have the same keys as both follow from TestWallets."
+              (expecs, vals) = unzip $ Map.elems $ theseToPair <$> zipMaps expected vs
+           in assertValues expecs vals
+
+    zipMaps :: Ord a => Map a b -> Map a c -> Map a (These b c)
+    zipMaps mb mc =
+      let f (This b) (That c) = These b c
+          f _ _ = error "All left are This and all right are That."
+       in Map.unionWith f (This <$> mb) (That <$> mc)
 
--- | Run a contract using the first wallet as own wallet, and return `ExecutionResult`.
+-- | Run a contract using the first wallet (in the order of how initializations are written) as own wallet, and return `ExecutionResult`.
 -- This could be used by itself, or combined with multiple other contracts.
 --
 -- @since 0.2
 withContract ::
   forall (w :: Type) (s :: Row Type) (e :: Type) (a :: Type).
   TestContractConstraints w e a =>
-  ([PaymentPubKeyHash] -> Contract w s e a) ->
+  (WalletLookups -> Contract w s e a) ->
   TestRunner w e a
-withContract = withContractAs 0
+withContract toContract = do
+  (_, wallets') <- ask
+  withContractAs (bwTag $ NonEmpty.head wallets') toContract
 
--- | Run a contract using the nth wallet as own wallet, and return `ExecutionResult`.
+-- | Run a contract using wallet with the given tag as own wallet, and return `ExecutionResult`.
 -- This could be used by itself, or combined with multiple other contracts.
 --
 -- @since 0.2
 withContractAs ::
   forall (w :: Type) (s :: Row Type) (e :: Type) (a :: Type).
   TestContractConstraints w e a =>
-  Int ->
-  ([PaymentPubKeyHash] -> Contract w s e a) ->
+  Text ->
+  (WalletLookups -> Contract w s e a) ->
   TestRunner w e a
-withContractAs walletIdx toContract = do
+withContractAs walletName toContract = do
   (cEnv, wallets') <- ask
   let -- pick wallet for Contract's "own PKH", other wallets PKHs will be provided
       -- to the user in `withContractAs`
-      (ownWallet, otherWallets) = separateWallets walletIdx wallets'
-
-      {- these are `PaymentPubKeyHash`es of all wallets used in test case
-      they stay in list is same order as `TestWallets` defined in test case
-      so collected Values will be in same order as well
-      it is important to preserve this order for Values check with `assertValues`
-      as there is no other mechanism atm to match `TestWallet` with collected `Value`
-      -}
-      collectValuesPkhs :: NonEmpty PaymentPubKeyHash
-      collectValuesPkhs = fmap ledgerPaymentPkh wallets'
-
-      -- wallet `PaymentPubKeyHash`es that will be available in
-      -- `withContract` and `withContractAs`
-      otherWalletsPkhs :: [PaymentPubKeyHash]
-      otherWalletsPkhs = fmap ledgerPaymentPkh otherWallets
+      (ownWallet, otherWallets) = separateWallets walletName $ NonEmpty.toList wallets'
+
+      -- without own wallet
+      otherLookups :: Map Text WalletInfo
+      otherLookups = lookupsMap otherWallets
+
+      -- to be passed to the user, without own wallet
+      walletLookups = makeWalletLookups otherLookups
+
+      -- these are `PaymentPubKeyHash`es of all wallets used in test case
+      collectValuesAddr = ownAddress <$> Map.insert (bwTag ownWallet) (makeWalletInfo ownWallet) otherLookups
 
       -- contract that gets all the values present at the test wallets.
-      valuesAtWallet :: Contract w s e (NonEmpty Value)
+      valuesAtWallet :: Contract w s e (Map Text Value)
       valuesAtWallet =
         void (waitNSlots 1)
-          >> traverse (valueAt . (`pubKeyHashAddress` Nothing)) collectValuesPkhs
+          >> traverse valueAt collectValuesAddr
 
   -- run the test contract
-  execRes <- liftIO $ runContract cEnv ownWallet (toContract otherWalletsPkhs)
+  execRes <- liftIO $ runContract cEnv ownWallet (toContract walletLookups)
 
   -- get all the values present at the test wallets after the user given contracts has been executed.
   execValues <- liftIO $ runContract cEnv ownWallet valuesAtWallet
@@ -347,12 +379,15 @@ withContractAs walletIdx toContract = do
     Left e -> fail $ "Failed to get values. Error: " ++ show e
     Right values -> return $ execRes {outcome = (,values) <$> outcome execRes}
   where
-    separateWallets :: forall b. Int -> NonEmpty b -> (b, [b])
-    separateWallets i xss
-      | (xs, y : ys) <- NonEmpty.splitAt i xss = (y, xs <> ys)
-      | otherwise = error $ "Should fail: bad wallet index for own wallet: " <> show i
+    separateWallets :: Text -> [BpiWallet] -> (BpiWallet, [BpiWallet])
+    separateWallets tag =
+      let p = (== tag) . bwTag
+          loop ys = \case
+            (a : xs) -> if p a then (a, xs <> ys) else loop (a : ys) xs
+            [] -> error $ "Should fail: bad wallet tag for own wallet: " <> show tag
+       in loop []
 
-newtype StatsReport w e a = StatsReport (IO (ExecutionResult w e (a, NonEmpty Value)))
+newtype StatsReport w e a = StatsReport (IO (ExecutionResult w e (a, Map Text Value)))
 
 instance
   forall (w :: Type) (e :: Type) (a :: Type).
@@ -369,7 +404,7 @@ instance
   testOptions = Tagged []
 
 -- | Test case used internally for logs printing.
-data LogsReport w e a = LogsReport LogsReportOption (IO (ExecutionResult w e (a, NonEmpty Value)))
+data LogsReport w e a = LogsReport LogsReportOption (IO (ExecutionResult w e (a, Map Text Value)))
 
 -- | TraceOption stripped to what LogsReport wants to know.
 data LogsReportOption
diff --git a/src/Test/Plutip/Contract/Init.hs b/src/Test/Plutip/Contract/Init.hs
index a963b852..0aaa470b 100644
--- a/src/Test/Plutip/Contract/Init.hs
+++ b/src/Test/Plutip/Contract/Init.hs
@@ -24,96 +24,100 @@ import Ledger.Value qualified as Value
 
 import Numeric.Positive (Positive)
 
-import Test.Plutip.Contract.Types (
-  TestWallet (TestWallet, twExpected, twInitDistribuition),
-  TestWallets (TestWallets, unTestWallets),
+import Test.Plutip.Internal.BotPlutusInterface.Run (defCollateralSize)
+import Test.Plutip.Internal.BotPlutusInterface.Types (
+  -- WalletSpec (wsExpected, wsInitDistribiution),
+  TestWallet (TestWallet),
+  TestWallets,
   ValueOrdering (VEq),
+  WalletTag,
+  mkWallet,
  )
-import Test.Plutip.Internal.BotPlutusInterface.Run (defCollateralSize)
 import Test.Plutip.Tools (ada)
 
 -- | Create a wallet with the given amounts of lovelace.
 --  Each amount will be sent to address as separate UTXO.
 --
 -- @since 0.2
-initLovelace :: [Positive] -> TestWallets
-initLovelace initial = TestWallets $ TestWallet initial Nothing :| []
+initLovelace :: WalletTag t -> [Positive] -> TestWallets
+initLovelace tag initial = mkWallet initial Nothing tag :| []
 
 -- | Create a wallet with the given amounts of lovelace, and after contract execution
 -- compare the values at the wallet address with the given ordering and value.
 --
 -- @since 0.2
-initLovelaceAssertValueWith :: [Positive] -> ValueOrdering -> Value -> TestWallets
-initLovelaceAssertValueWith initial ord expect = TestWallets $ TestWallet initial (Just (ord, expect)) :| []
+initLovelaceAssertValueWith :: WalletTag t -> [Positive] -> ValueOrdering -> Value -> TestWallets
+initLovelaceAssertValueWith tag initial ord expect = mkWallet initial (Just (ord, expect)) tag :| []
 
 -- | Create a wallet with the given amounts of lovelace, and after contract execution
 -- check if values at the wallet address are equal to a given value.
 --
 -- @since 0.2
-initLovelaceAssertValue :: [Positive] -> Value -> TestWallets
-initLovelaceAssertValue initial = initLovelaceAssertValueWith initial VEq
+initLovelaceAssertValue :: WalletTag t -> [Positive] -> Value -> TestWallets
+initLovelaceAssertValue tag initial = initLovelaceAssertValueWith tag initial VEq
 
 -- | Create a wallet with the given amounts of lovelace, and after contract execution
 -- compare the values at the wallet address with the given ordering and lovelace amount.
 --
 -- @since 0.2
-initAndAssertLovelaceWith :: [Positive] -> ValueOrdering -> Positive -> TestWallets
-initAndAssertLovelaceWith initial ord expect =
-  initLovelaceAssertValueWith initial ord (Ada.lovelaceValueOf (fromIntegral expect))
+initAndAssertLovelaceWith :: WalletTag t -> [Positive] -> ValueOrdering -> Positive -> TestWallets
+initAndAssertLovelaceWith tag initial ord expect =
+  initLovelaceAssertValueWith tag initial ord (Ada.lovelaceValueOf (fromIntegral expect))
 
 -- | Create a wallet with the given amounts of lovelace, and after contract execution
 -- check if values at the wallet address are equal to a given lovelace amount.
 --
 -- @since 0.2
-initAndAssertLovelace :: [Positive] -> Positive -> TestWallets
-initAndAssertLovelace initial expect =
-  initLovelaceAssertValue initial (Ada.lovelaceValueOf (fromIntegral expect))
+initAndAssertLovelace :: WalletTag t -> [Positive] -> Positive -> TestWallets
+initAndAssertLovelace tag initial expect =
+  initLovelaceAssertValue tag initial (Ada.lovelaceValueOf (fromIntegral expect))
 
 -- | Create a wallet with the given amounts of Ada.
 --
 -- @since 0.2
-initAda :: [Positive] -> TestWallets
-initAda initial = initLovelace (map ada initial)
+initAda :: WalletTag t -> [Positive] -> TestWallets
+initAda tag initial = initLovelace tag (map ada initial)
 
 -- | Create a wallet with the given amounts of Ada, and after contract execution
 -- compare the values at the wallet address with the given ordering and value.
 --
 -- @since 0.2
-initAdaAssertValueWith :: [Positive] -> ValueOrdering -> Value -> TestWallets
-initAdaAssertValueWith initial = initLovelaceAssertValueWith (map ada initial)
+initAdaAssertValueWith :: WalletTag t -> [Positive] -> ValueOrdering -> Value -> TestWallets
+initAdaAssertValueWith tag initial = initLovelaceAssertValueWith tag (map ada initial)
 
 -- | Create a wallet with the given amounts of Ada, and after contract execution
 -- check if values at the wallet address are equal to a given value.
 --
 -- @since 0.2
-initAdaAssertValue :: [Positive] -> Value -> TestWallets
-initAdaAssertValue initial = initLovelaceAssertValue (map ada initial)
+initAdaAssertValue :: WalletTag t -> [Positive] -> Value -> TestWallets
+initAdaAssertValue tag initial = initLovelaceAssertValue tag (map ada initial)
 
 -- | Create a wallet with the given amounts of Ada, and after contract execution
 -- compare the values at the wallet address with the given ordering and ada amount.
 --
 -- @since 0.2
-initAndAssertAdaWith :: [Positive] -> ValueOrdering -> Positive -> TestWallets
-initAndAssertAdaWith initial ord expect =
-  initAndAssertLovelaceWith (map ada initial) ord (ada expect)
+initAndAssertAdaWith :: WalletTag t -> [Positive] -> ValueOrdering -> Positive -> TestWallets
+initAndAssertAdaWith tag initial ord expect =
+  initAndAssertLovelaceWith tag (map ada initial) ord (ada expect)
 
 -- | Create a wallet with the given amounts of Ada, and after contract execution
 -- check if values at the wallet address are equal to a given ada amount.
 --
 -- @since 0.2
-initAndAssertAda :: [Positive] -> Positive -> TestWallets
-initAndAssertAda initial expect =
-  initAndAssertLovelace (map ada initial) (ada expect)
+initAndAssertAda :: WalletTag t -> [Positive] -> Positive -> TestWallets
+initAndAssertAda tag initial expect =
+  initAndAssertLovelace tag (map ada initial) (ada expect)
 
 -- | Initialize all the 'TestWallets' with the collateral utxo and
 --   adjust the 'twExpected' value accordingly.
 withCollateral :: TestWallets -> TestWallets
-withCollateral TestWallets {..} = TestWallets $ NonEmpty.map go unTestWallets
+withCollateral = NonEmpty.map go
   where
     go :: TestWallet -> TestWallet
-    go TestWallet {..} =
+    go (TestWallet tag dist expected) =
       TestWallet
-        { twInitDistribuition = fromInteger defCollateralSize : twInitDistribuition
-        , twExpected =
-            second (Value.unionWith (+) $ Ada.lovelaceValueOf defCollateralSize) <$> twExpected
-        }
+        tag
+        (fromInteger defCollateralSize : dist)
+        (second (Value.unionWith (+) collateral) <$> expected)
+
+    collateral = Ada.lovelaceValueOf defCollateralSize
diff --git a/src/Test/Plutip/Contract/Types.hs b/src/Test/Plutip/Contract/Types.hs
index 1fd20741..2f63f3f7 100644
--- a/src/Test/Plutip/Contract/Types.hs
+++ b/src/Test/Plutip/Contract/Types.hs
@@ -1,24 +1,22 @@
 {-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE GADTs #-}
 
 module Test.Plutip.Contract.Types (
   TestContractConstraints,
   TestContract (..),
-  TestWallets (TestWallets, unTestWallets),
-  TestWallet (..),
-  compareValuesWith,
-  ValueOrdering (..),
+  WalletTag (..),
 ) where
 
 import Data.Aeson (ToJSON)
 import Data.Bool (bool)
 import Data.Dynamic (Typeable)
 import Data.Kind (Type)
-import Data.List.NonEmpty (NonEmpty)
+import Data.Map (Map)
 import Data.Tagged (Tagged (Tagged))
+import Data.Text (Text)
 import Ledger.Value (Value)
-import Ledger.Value qualified as Value
-import Numeric.Positive (Positive)
 import Plutus.Contract (AsContractError)
+import Test.Plutip.Internal.BotPlutusInterface.Types (WalletTag (BaseTag, EntTag))
 import Test.Plutip.Internal.Types (
   ExecutionResult,
  )
@@ -42,7 +40,7 @@ data TestContract (w :: Type) (e :: Type) (a :: Type)
   = TestContract
       (Predicate w e a)
       -- ^ Info about check to perform and how to report results
-      (IO (ExecutionResult w e (a, NonEmpty Value)))
+      (IO (ExecutionResult w e (a, Map Text Value)))
       -- ^ Result of contract execution
   deriving stock (Typeable)
 
@@ -64,21 +62,3 @@ instance
         (pCheck predicate result)
 
   testOptions = Tagged []
-
-newtype TestWallets = TestWallets {unTestWallets :: NonEmpty TestWallet}
-  deriving newtype (Semigroup)
-
-data TestWallet = TestWallet
-  { twInitDistribuition :: [Positive]
-  , twExpected :: Maybe (ValueOrdering, Value)
-  }
-
-data ValueOrdering = VEq | VGt | VLt | VGEq | VLEq
-
--- | Value doesn't have an Ord instance, so we cannot use `compare`
-compareValuesWith :: ValueOrdering -> Value -> Value -> Bool
-compareValuesWith VEq = (==)
-compareValuesWith VGt = Value.gt
-compareValuesWith VLt = Value.lt
-compareValuesWith VGEq = Value.geq
-compareValuesWith VLEq = Value.leq
diff --git a/src/Test/Plutip/Contract/Values.hs b/src/Test/Plutip/Contract/Values.hs
index 4fb01149..58af66c7 100644
--- a/src/Test/Plutip/Contract/Values.hs
+++ b/src/Test/Plutip/Contract/Values.hs
@@ -10,8 +10,6 @@ import Data.Aeson.Extras (encodeByteString)
 import Data.Either (fromRight)
 import Data.Kind (Type)
 import Data.List (find)
-import Data.List.NonEmpty (NonEmpty)
-import Data.List.NonEmpty qualified as NonEmpty
 import Data.Map qualified as Map
 import Data.Row (Row)
 import Data.Text (Text)
@@ -24,7 +22,7 @@ import Ledger.Value qualified as Value
 import Plutus.Contract (AsContractError, Contract, utxosAt)
 import PlutusTx.Builtins (fromBuiltin)
 
-import Test.Plutip.Contract.Types (
+import Test.Plutip.Internal.BotPlutusInterface.Types (
   ValueOrdering (VEq, VGEq, VGt, VLEq, VLt),
   compareValuesWith,
  )
@@ -42,11 +40,11 @@ valueAt addr = do
     utxoValue (PublicKeyChainIndexTxOut _ v _ _) = v
     utxoValue (ScriptChainIndexTxOut _ v _ _ _) = v
 
-assertValues :: NonEmpty (Maybe (ValueOrdering, Value)) -> NonEmpty Value -> Either Text ()
+assertValues :: [Maybe (ValueOrdering, Value)] -> [Value] -> Either Text ()
 assertValues expected values =
   maybe (Right ()) (Left . report) $
     find findFailing $
-      zip3 [0 :: Int ..] (NonEmpty.toList expected) (NonEmpty.toList values)
+      zip3 [0 :: Int ..] expected values
   where
     findFailing (_, Nothing, _) = False
     findFailing (_, Just (ord, v), v') = not (compareValuesWith ord v' v)
diff --git a/src/Test/Plutip/Internal/BotPlutusInterface/Keys.hs b/src/Test/Plutip/Internal/BotPlutusInterface/Keys.hs
index 39096e1c..b3eea7a0 100644
--- a/src/Test/Plutip/Internal/BotPlutusInterface/Keys.hs
+++ b/src/Test/Plutip/Internal/BotPlutusInterface/Keys.hs
@@ -1,43 +1,89 @@
-module Test.Plutip.Internal.BotPlutusInterface.Keys (genKeyPair, genKeyPairs) where
+{-# LANGUAGE OverloadedStrings #-}
 
-import Cardano.Api (AsType (AsPaymentKey), Key (VerificationKey, getVerificationKey, verificationKeyHash), PaymentKey, SigningKey, TextEnvelopeDescr, generateSigningKey, writeFileTextEnvelope)
+module Test.Plutip.Internal.BotPlutusInterface.Keys (
+  KeyPair (sKey, vKey),
+  StakeKeyPair (sSKey, sVKey),
+  genKeyPair,
+  writeKeyPair,
+  genStakeKeyPair,
+  writeStakeKeyPairs,
+  signingKeyFilePathInDir,
+  verificationKeyFilePathInDir,
+  stakingSigningKeyFilePathInDir,
+  stakingVerificationKeyFilePathInDir,
+) where
+
+import Cardano.Api (writeFileTextEnvelope)
+import Cardano.Api qualified as CAPI
+import Data.Text qualified as Text
 import System.FilePath ((<.>), (</>))
 
 data KeyPair = KeyPair
-  { sKey :: SigningKey PaymentKey
-  , vKey :: VerificationKey PaymentKey
+  { sKey :: CAPI.SigningKey CAPI.PaymentKey
+  , vKey :: CAPI.VerificationKey CAPI.PaymentKey
   }
   deriving stock (Show)
 
-genKeyPair :: IO KeyPair
-genKeyPair = do
-  sKey <- generateSigningKey AsPaymentKey
-  return $ KeyPair sKey (getVerificationKey sKey)
+data StakeKeyPair = StakeKeyPair
+  { sSKey :: CAPI.SigningKey CAPI.StakeKey
+  , sVKey :: CAPI.VerificationKey CAPI.StakeKey
+  }
+  deriving stock (Show)
+
+pSKeyDesc, pVKeyDesc, sSKeyDesc, sVKeyDesc :: CAPI.TextEnvelopeDescr
+pSKeyDesc = "Payment Signing Key"
+pVKeyDesc = "Payment Verification Key"
+sSKeyDesc = "Stake Signing Key"
+sVKeyDesc = "Stake Verification Key"
 
 -- | Helper to generate key pairs.
 -- Can be further developed to generate test keys for test wallets
 -- to work with `bot-plutus-interface`
--- >>> genKeyPairs "cluster-data/known_wallets" "signing-key-" "verification-key-"
-genKeyPairs :: FilePath -> String -> String -> IO ()
-genKeyPairs outDir sKeyPrefix vKeyPrefix = do
-  sKey <- generateSigningKey AsPaymentKey
-  let skeyDesc, vkeyDesc :: TextEnvelopeDescr
-      skeyDesc = "Payment Signing Key"
-      vkeyDesc = "Payment Verification Key"
-
-      vKey = getVerificationKey sKey
-      hash = verificationKeyHash vKey
-
-      skeyPath = rmQuotes $ outDir </> sKeyPrefix ++ showHash hash <.> "skey"
-      vkeyPath = rmQuotes $ outDir </> vKeyPrefix ++ showHash hash <.> "vkey"
-
-      showHash = rmQuotes . show
-  res <-
-    sequence
-      [ writeFileTextEnvelope skeyPath (Just skeyDesc) sKey
-      , writeFileTextEnvelope vkeyPath (Just vkeyDesc) vKey
-      ]
-  print res
-
-rmQuotes :: String -> String
-rmQuotes = filter (/= '"')
+genKeyPair :: IO KeyPair
+genKeyPair = do
+  sKey <- CAPI.generateSigningKey CAPI.AsPaymentKey
+  return $ KeyPair sKey (CAPI.getVerificationKey sKey)
+
+writeKeyPair :: FilePath -> KeyPair -> IO [Either (CAPI.FileError ()) ()]
+writeKeyPair outDir keyPair = do
+  let hash = CAPI.verificationKeyHash $ vKey keyPair
+
+      skeyPath = signingKeyFilePathInDir outDir hash
+      vkeyPath = verificationKeyFilePathInDir outDir hash
+
+  sequence
+    [ writeFileTextEnvelope skeyPath (Just pSKeyDesc) (sKey keyPair)
+    , writeFileTextEnvelope vkeyPath (Just pVKeyDesc) (vKey keyPair)
+    ]
+
+genStakeKeyPair :: IO StakeKeyPair
+genStakeKeyPair = do
+  sKey <- CAPI.generateSigningKey CAPI.AsStakeKey
+  return $ StakeKeyPair sKey (CAPI.getVerificationKey sKey)
+
+writeStakeKeyPairs :: FilePath -> StakeKeyPair -> IO [Either (CAPI.FileError ()) ()]
+writeStakeKeyPairs dir stakeKeyPair = do
+  let hash = CAPI.verificationKeyHash $ sVKey stakeKeyPair
+
+      skeyPath = stakingSigningKeyFilePathInDir dir hash
+      vkeyPath = stakingVerificationKeyFilePathInDir dir hash
+
+  sequence
+    [ writeFileTextEnvelope skeyPath (Just sSKeyDesc) (sSKey stakeKeyPair)
+    , writeFileTextEnvelope vkeyPath (Just sVKeyDesc) (sVKey stakeKeyPair)
+    ]
+
+signingKeyFilePathInDir :: FilePath -> CAPI.Hash CAPI.PaymentKey -> FilePath
+signingKeyFilePathInDir dir vkh = keyFilePathInDir dir "signing-key" vkh "skey"
+
+verificationKeyFilePathInDir :: FilePath -> CAPI.Hash CAPI.PaymentKey -> FilePath
+verificationKeyFilePathInDir dir vkh = keyFilePathInDir dir "verification-key" vkh "vkey"
+
+stakingSigningKeyFilePathInDir :: FilePath -> CAPI.Hash CAPI.StakeKey -> FilePath
+stakingSigningKeyFilePathInDir dir vkh = keyFilePathInDir dir "staking-signing-key" vkh "skey"
+
+stakingVerificationKeyFilePathInDir :: FilePath -> CAPI.Hash CAPI.StakeKey -> FilePath
+stakingVerificationKeyFilePathInDir dir vkh = keyFilePathInDir dir "staking-verification-key" vkh "vkey"
+
+keyFilePathInDir :: forall a. CAPI.SerialiseAsRawBytes (CAPI.Hash a) => FilePath -> String -> CAPI.Hash a -> String -> FilePath
+keyFilePathInDir dir pref h ext = dir </> pref <> "-" <> Text.unpack (CAPI.serialiseToRawBytesHexText h) <.> ext
diff --git a/src/Test/Plutip/Internal/BotPlutusInterface/Lookups.hs b/src/Test/Plutip/Internal/BotPlutusInterface/Lookups.hs
new file mode 100644
index 00000000..1a96b8ba
--- /dev/null
+++ b/src/Test/Plutip/Internal/BotPlutusInterface/Lookups.hs
@@ -0,0 +1,88 @@
+module Test.Plutip.Internal.BotPlutusInterface.Lookups (
+  WalletLookups (lookupAddress, lookupWallet),
+  makeWalletInfo,
+  makeWalletLookups,
+  lookupsMap,
+) where
+
+import Control.Lens (withPrism)
+import Control.Monad.Except (MonadError (throwError))
+import Data.Kind (Type)
+import Data.Map (Map)
+import Data.Map qualified as Map
+import Data.Row (Row)
+import Data.Text (Text)
+import Ledger (Address)
+import Plutus.Contract (Contract, ContractError (OtherContractError))
+import Plutus.Contract.Error (AsContractError, _ContractError)
+import Test.Plutip.Internal.BotPlutusInterface.Types (
+  BaseWallet (BaseWallet),
+  BpiWallet (bwTag),
+  EntWallet (EntWallet),
+  WalletInfo,
+  WalletTag (BaseTag, EntTag),
+  ownAddress,
+ )
+import Test.Plutip.Internal.BotPlutusInterface.Wallet (walletPaymentPkh, walletStakePkh)
+
+-- Error messages for wallet lookup fails.
+expectedEnterpriseWallet, expectedWalletWithStakeKeys :: Text
+expectedEnterpriseWallet = "Expected base address wallet, got one with staking keys."
+expectedWalletWithStakeKeys = "Expected base address wallet, got one with staking keys."
+badWalletTag :: Text -> Text
+badWalletTag tag = "Wallet not found by tag '" <> tag <> "'."
+
+-- | Type to be used for looking up wallet informations. Wallets accessed by their k typed names.
+data WalletLookups = WalletLookups
+  { lookupWallet ::
+      forall (t :: Type) (w :: Type) (s :: Row Type) (e :: Type).
+      AsContractError e =>
+      WalletTag t ->
+      Contract w s e t
+  , lookupAddress ::
+      forall (w :: Type) (s :: Row Type) (e :: Type).
+      AsContractError e =>
+      Text ->
+      Contract w s e Address
+  }
+
+makeWalletInfo :: BpiWallet -> WalletInfo
+makeWalletInfo w =
+  maybe
+    (Right $ EntWallet (walletPaymentPkh w))
+    (Left . BaseWallet (walletPaymentPkh w))
+    (walletStakePkh w)
+
+lookupsMap :: [BpiWallet] -> Map Text WalletInfo
+lookupsMap bpiWalls =
+  Map.fromList $
+    (\w -> (bwTag w, makeWalletInfo w)) <$> bpiWalls
+
+makeWalletLookups ::
+  Map Text WalletInfo ->
+  WalletLookups
+makeWalletLookups lookups =
+  WalletLookups
+    { lookupWallet = lookupTaggedWallet lookups
+    , lookupAddress = \tag ->
+        maybe (toError $ badWalletTag tag) pure $
+          Map.lookup tag $ ownAddress <$> lookups
+    }
+  where
+    toError :: MonadError e m => AsContractError e => Text -> m a
+    toError = throwError . (\e -> withPrism _ContractError $ \f _ -> f e) . OtherContractError
+
+    lookupTaggedWallet ::
+      forall (w :: Type) (s :: Row Type) (e :: Type) (t :: Type).
+      (AsContractError e) =>
+      Map Text WalletInfo ->
+      WalletTag t ->
+      Contract w s e t
+    lookupTaggedWallet ws (EntTag tag) = case Map.lookup tag ws of
+      Nothing -> toError $ badWalletTag tag
+      Just (Right res@(EntWallet _)) -> pure res
+      Just (Left (BaseWallet _ _)) -> toError expectedEnterpriseWallet
+    lookupTaggedWallet ws (BaseTag tag) = case Map.lookup tag ws of
+      Nothing -> toError $ badWalletTag tag
+      Just (Right (EntWallet _)) -> toError expectedWalletWithStakeKeys
+      Just (Left res@(BaseWallet _ _)) -> pure res
diff --git a/src/Test/Plutip/Internal/BotPlutusInterface/Run.hs b/src/Test/Plutip/Internal/BotPlutusInterface/Run.hs
index 94731345..db87cbd4 100644
--- a/src/Test/Plutip/Internal/BotPlutusInterface/Run.hs
+++ b/src/Test/Plutip/Internal/BotPlutusInterface/Run.hs
@@ -57,11 +57,13 @@ import Data.Kind (Type)
 import Data.Row (Row)
 import Data.Text qualified as Text
 import Data.UUID.V4 qualified as UUID
+import Ledger (unPaymentPubKeyHash)
 import Plutus.Contract (Contract)
 import Plutus.PAB.Core.ContractInstance.STM (Activity (Active))
 import Test.Plutip.Config (PlutipConfig (budgetMultiplier))
+import Test.Plutip.Internal.BotPlutusInterface.Lookups (makeWalletInfo)
 import Test.Plutip.Internal.BotPlutusInterface.Setup qualified as BIS
-import Test.Plutip.Internal.BotPlutusInterface.Wallet (BpiWallet (walletPkh))
+import Test.Plutip.Internal.BotPlutusInterface.Types (BpiWallet, ownPaymentPubKeyHash, ownStakePubKeyHash)
 import Test.Plutip.Internal.Types (
   ClusterEnv (chainIndexUrl, networkId, plutipConf),
   ExecutionResult (ExecutionResult),
@@ -133,8 +135,8 @@ runContractWithLogLvl logLvl cEnv bpiWallet contract = do
         , pcDryRun = False
         , pcProtocolParamsFile = Text.pack $ BIS.pParamsFile cEnv
         , pcLogLevel = logLvl
-        , pcOwnPubKeyHash = walletPkh bpiWallet
-        , pcOwnStakePubKeyHash = Nothing
+        , pcOwnPubKeyHash = unPaymentPubKeyHash $ ownPaymentPubKeyHash walletInfo
+        , pcOwnStakePubKeyHash = ownStakePubKeyHash walletInfo
         , pcTipPollingInterval = 1_000_000
         , pcPort = 9080
         , pcEnableTxEndpoint = False
@@ -146,6 +148,8 @@ runContractWithLogLvl logLvl cEnv bpiWallet contract = do
         , pcCollateralSize = fromInteger defCollateralSize
         }
 
+    walletInfo = makeWalletInfo bpiWallet
+
     runContract' :: ContractEnvironment w -> m (ExecutionResult w e a)
     runContract' contractEnv = do
       res <- liftIO $ try @SomeException (BIC.runContract contractEnv contract)
diff --git a/src/Test/Plutip/Internal/BotPlutusInterface/Setup.hs b/src/Test/Plutip/Internal/BotPlutusInterface/Setup.hs
index dc1339de..4ebab9f1 100644
--- a/src/Test/Plutip/Internal/BotPlutusInterface/Setup.hs
+++ b/src/Test/Plutip/Internal/BotPlutusInterface/Setup.hs
@@ -13,12 +13,11 @@ import Cardano.Api qualified as CAPI
 import Cardano.Launcher.Node (nodeSocketFile)
 import Data.Aeson (encodeFile)
 import Data.Foldable (traverse_)
-import Plutus.V1.Ledger.Api (PubKeyHash (PubKeyHash))
-import PlutusTx.Builtins qualified as PlutusTx
 import System.Directory (createDirectoryIfMissing, doesDirectoryExist)
 import System.Environment (setEnv)
 import System.FilePath ((</>))
 import Test.Plutip.Config (PlutipConfig (extraSigners))
+import Test.Plutip.Internal.BotPlutusInterface.Keys (signingKeyFilePathInDir)
 import Test.Plutip.Internal.Types (ClusterEnv (plutipConf, supportDir), nodeSocket)
 import Test.Plutip.Tools.CardanoApi (queryProtocolParams)
 
@@ -70,12 +69,10 @@ runSetup cEnv = do
           Left fileError -> error $ displayError fileError
           Right sKey -> addExtraSigner $ Right sKey
       (Right sKey) -> do
-        let vKey = CAPI.getVerificationKey sKey
-            pkh = PubKeyHash . PlutusTx.toBuiltin . CAPI.serialiseToRawBytes $ CAPI.verificationKeyHash vKey
-            keyFilename = "signing-key-" <> show pkh <> ".skey"
+        let vKeyHash = CAPI.verificationKeyHash $ CAPI.getVerificationKey sKey
         g <-
           CAPI.writeFileTextEnvelope
-            (keysDir cEnv </> keyFilename)
+            (signingKeyFilePathInDir (keysDir cEnv) vKeyHash)
             Nothing
             sKey
         case g of
diff --git a/src/Test/Plutip/Internal/BotPlutusInterface/Types.hs b/src/Test/Plutip/Internal/BotPlutusInterface/Types.hs
index f0e6aa3a..8928ae07 100644
--- a/src/Test/Plutip/Internal/BotPlutusInterface/Types.hs
+++ b/src/Test/Plutip/Internal/BotPlutusInterface/Types.hs
@@ -1,8 +1,107 @@
+{-# LANGUAGE GADTs #-}
+
 module Test.Plutip.Internal.BotPlutusInterface.Types (
   BpiError (..),
+  BpiWallet (BpiWallet, payKeys, stakeKeys, bwTag),
+  TestWallets,
+  ValueOrdering (VEq, VGt, VLt, VGEq, VLEq),
+  compareValuesWith,
+  WalletTag (..),
+  TestWallet (..),
+  WalletInfo,
+  ownPaymentPubKeyHash,
+  ownStakePubKeyHash,
+  ownAddress,
+  getTag,
+  mkWallet,
+  BaseWallet (..),
+  EntWallet (..),
 ) where
 
+import Data.Data (Typeable)
+import Data.List.NonEmpty (NonEmpty)
+import Data.Text (Text)
+import Ledger (Address, PaymentPubKeyHash, StakePubKeyHash, Value, pubKeyHashAddress)
+import Ledger.Value qualified as Value
+import Numeric.Positive (Positive)
+import Test.Plutip.Internal.BotPlutusInterface.Keys (KeyPair, StakeKeyPair)
+
+-- | Tag of the wallet that gives wallet name and specifies type of address it will have:
+-- base or enterprise.
+-- Also used in lookups and specifies expected type of returned wallet.
+--
+-- `t` type parameter is the type of wallet that will be accessible from `WalletLookups`.
+data WalletTag t where
+  -- | Option to create wallet with base address: has both payment and staking keys
+  BaseTag :: Text -> WalletTag BaseWallet
+  -- | Option to create wallet with enterprise address: has only payment keys
+  EntTag :: Text -> WalletTag EntWallet
+
+deriving stock instance Show (WalletTag t)
+deriving stock instance Eq (WalletTag t)
+
 data BpiError
   = SignKeySaveError !String
   | BotInterfaceDirMissing
+  deriving stock (Show, Eq)
+
+-- | Wallet that can be used by bot interface,
+--  backed by `.skey` file when added to cluster with `addSomeWallet`
+data BpiWallet = BpiWallet
+  { payKeys :: KeyPair
+  , stakeKeys :: Maybe StakeKeyPair
+  , bwTag :: Text
+  }
   deriving stock (Show)
+
+type TestWallets = NonEmpty TestWallet
+
+data TestWallet = forall t.
+  TestWallet
+  { twTag :: WalletTag t
+  , twDistribution :: [Positive]
+  , twExpected :: Maybe (ValueOrdering, Value)
+  }
+
+-- | Make TestWallet, takes utxo distribution, value assertions and WalletTag as arguments.
+mkWallet :: [Positive] -> Maybe (ValueOrdering, Value) -> WalletTag t -> TestWallet
+mkWallet initDistribiution expected tag =
+  TestWallet tag initDistribiution expected
+
+getTag :: TestWallet -> Text
+getTag (TestWallet tag _ _) = getTag' tag
+  where
+    getTag' :: WalletTag t -> Text
+    getTag' = \case
+      BaseTag tag' -> tag'
+      EntTag tag' -> tag'
+
+data ValueOrdering = VEq | VGt | VLt | VGEq | VLEq
+
+-- | Value doesn't have an Ord instance, so we cannot use `compare`
+compareValuesWith :: ValueOrdering -> Value -> Value -> Bool
+compareValuesWith VEq = (==)
+compareValuesWith VGt = Value.gt
+compareValuesWith VLt = Value.lt
+compareValuesWith VGEq = Value.geq
+compareValuesWith VLEq = Value.leq
+
+-- | Type holding wallet information as seen with wallet lookups. Used internally only.
+type WalletInfo = Either BaseWallet EntWallet
+
+-- | Base address wallet: supported by both Payment and Staking keys.
+data BaseWallet = BaseWallet {getBasePkh :: PaymentPubKeyHash, getSpkh :: StakePubKeyHash}
+  deriving stock (Eq, Show, Typeable)
+
+-- | Enterprise address wallet: supported only by Payment keys.
+newtype EntWallet = EntWallet {getPkh :: PaymentPubKeyHash}
+  deriving stock (Eq, Show, Typeable)
+
+ownPaymentPubKeyHash :: WalletInfo -> PaymentPubKeyHash
+ownPaymentPubKeyHash = either getBasePkh getPkh
+
+ownStakePubKeyHash :: WalletInfo -> Maybe StakePubKeyHash
+ownStakePubKeyHash = either (Just . getSpkh) (const Nothing)
+
+ownAddress :: WalletInfo -> Address
+ownAddress w = pubKeyHashAddress (ownPaymentPubKeyHash w) (ownStakePubKeyHash w)
diff --git a/src/Test/Plutip/Internal/BotPlutusInterface/Wallet.hs b/src/Test/Plutip/Internal/BotPlutusInterface/Wallet.hs
index f7f38285..08ef73b7 100644
--- a/src/Test/Plutip/Internal/BotPlutusInterface/Wallet.hs
+++ b/src/Test/Plutip/Internal/BotPlutusInterface/Wallet.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE TupleSections #-}
+
 module Test.Plutip.Internal.BotPlutusInterface.Wallet (
   BpiWallet (..),
   addSomeWallet,
@@ -6,12 +8,16 @@ module Test.Plutip.Internal.BotPlutusInterface.Wallet (
   eitherAddSomeWalletDir,
   mkMainnetAddress,
   cardanoMainnetAddress,
-  ledgerPaymentPkh,
+  walletPaymentPkh,
+  walletStakePkh,
 ) where
 
-import Cardano.Api (AddressAny, PaymentKey, SigningKey, VerificationKey)
+import Cardano.Api (AddressAny)
 import Cardano.Api qualified as CAPI
+import Cardano.Api.Shelley qualified as CAPI
 import Cardano.BM.Data.Tracer (nullTracer)
+import Cardano.Ledger.BaseTypes as Shelley (Network (Mainnet))
+import Cardano.Ledger.Credential qualified as Shelley
 import Cardano.Wallet.Primitive.Types.Coin (Coin (Coin))
 import Cardano.Wallet.Shelley.Launch.Cluster (
   sendFaucetFundsTo,
@@ -20,29 +26,36 @@ import Control.Arrow (ArrowChoice (left))
 import Control.Monad (void)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Control.Monad.Reader (ReaderT, ask)
-import Data.Aeson.Extras (encodeByteString)
 import Data.Bool (bool)
+import Data.Either (isRight)
+import Data.Maybe (listToMaybe)
 import Data.Text qualified as Text
-import Ledger (PaymentPubKeyHash (PaymentPubKeyHash), PubKeyHash (PubKeyHash))
+import Ledger (
+  PaymentPubKeyHash (PaymentPubKeyHash),
+  PubKeyHash (PubKeyHash),
+  StakePubKeyHash (StakePubKeyHash),
+ )
 import Numeric.Positive (Positive)
-import Plutus.V1.Ledger.Api qualified as LAPI
-import PlutusTx.Builtins (fromBuiltin, toBuiltin)
+import PlutusTx.Builtins (toBuiltin)
 import System.Directory (createDirectoryIfMissing)
-import System.FilePath ((<.>), (</>))
+import Test.Plutip.Internal.BotPlutusInterface.Keys (
+  KeyPair (vKey),
+  StakeKeyPair (sVKey),
+  genKeyPair,
+  genStakeKeyPair,
+  writeKeyPair,
+  writeStakeKeyPairs,
+ )
 import Test.Plutip.Internal.BotPlutusInterface.Setup qualified as Setup
-import Test.Plutip.Internal.BotPlutusInterface.Types (BpiError (BotInterfaceDirMissing, SignKeySaveError))
+import Test.Plutip.Internal.BotPlutusInterface.Types (
+  BpiError (BotInterfaceDirMissing, SignKeySaveError),
+  BpiWallet (BpiWallet),
+  WalletTag (BaseTag, EntTag),
+  payKeys,
+  stakeKeys,
+ )
 import Test.Plutip.Internal.Types (ClusterEnv, nodeSocket, supportDir)
 
--- | Wallet that can be used by bot interface,
---  backed by `.skey` file when added to cluster with `addSomeWallet`
-data BpiWallet = BpiWallet
-  { walletPkh :: !PubKeyHash
-  , vrfKey :: VerificationKey PaymentKey
-  , signKey :: SigningKey PaymentKey
-  -- todo: do we need something else?
-  }
-  deriving stock (Show)
-
 {-  Add wallet with arbitrary address and specified amount of Ada.
   Each value specified in funds will be sent as separate UTXO.
 
@@ -50,14 +63,14 @@ During wallet addition `.skey` file with required name generated and saved
  to be used by bot interface.
  Directory for files could be obtained with `Test.Plutip.BotPlutusInterface.Setup.keysDir`
 -}
-eitherAddSomeWallet :: MonadIO m => [Positive] -> ReaderT ClusterEnv m (Either BpiError BpiWallet)
-eitherAddSomeWallet funds = eitherAddSomeWalletDir funds Nothing
+eitherAddSomeWallet :: MonadIO m => WalletTag t -> [Positive] -> ReaderT ClusterEnv m (Either BpiError BpiWallet)
+eitherAddSomeWallet tag funds = eitherAddSomeWalletDir tag funds Nothing
 
 -- | The same as `eitherAddSomeWallet`, but also
 -- saves the key file to a separate directory.
-eitherAddSomeWalletDir :: MonadIO m => [Positive] -> Maybe FilePath -> ReaderT ClusterEnv m (Either BpiError BpiWallet)
-eitherAddSomeWalletDir funds wallDir = do
-  bpiWallet <- createWallet
+eitherAddSomeWalletDir :: MonadIO m => WalletTag t -> [Positive] -> Maybe FilePath -> ReaderT ClusterEnv m (Either BpiError BpiWallet)
+eitherAddSomeWalletDir tag funds wallDir = do
+  bpiWallet <- createWallet tag
   saveWallets bpiWallet wallDir
     >>= \case
       Right _ -> sendFunds bpiWallet >> pure (Right bpiWallet)
@@ -76,27 +89,23 @@ eitherAddSomeWalletDir funds wallDir = do
 
 -- | Add wallet with arbitrary address and specified amount of Ada.
 -- (version of `eitherAddSomeWallet` that will throw an error in case of failure)
-addSomeWallet :: MonadIO m => [Positive] -> ReaderT ClusterEnv m BpiWallet
-addSomeWallet funds =
-  eitherAddSomeWallet funds >>= either (error . show) pure
+addSomeWallet :: MonadIO m => WalletTag t -> [Positive] -> ReaderT ClusterEnv m BpiWallet
+addSomeWallet tag funds =
+  eitherAddSomeWallet tag funds >>= either (error . show) pure
 
 -- | Version of `addSomeWallet` that also writes the
 -- wallet key file to a separate directory
-addSomeWalletDir :: MonadIO m => [Positive] -> Maybe FilePath -> ReaderT ClusterEnv m BpiWallet
-addSomeWalletDir funds wallDir =
-  eitherAddSomeWalletDir funds wallDir >>= either (error . show) pure
+addSomeWalletDir :: MonadIO m => WalletTag t -> [Positive] -> Maybe FilePath -> ReaderT ClusterEnv m BpiWallet
+addSomeWalletDir tag funds wallDir =
+  eitherAddSomeWalletDir tag funds wallDir >>= either (error . show) pure
 
-createWallet :: MonadIO m => m BpiWallet
-createWallet = do
-  sKey <- liftIO $ CAPI.generateSigningKey CAPI.AsPaymentKey
-  let vKey = CAPI.getVerificationKey sKey
-  return $ BpiWallet (toPkh vKey) vKey sKey
-  where
-    toPkh =
-      PubKeyHash
-        . toBuiltin
-        . CAPI.serialiseToRawBytes
-        . CAPI.verificationKeyHash
+createWallet :: MonadIO m => WalletTag t -> m BpiWallet
+createWallet tag = do
+  kp <- liftIO genKeyPair
+  (skp, k) <- case tag of
+    EntTag k -> pure (Nothing, k)
+    BaseTag k -> fmap ((,k) . Just) (liftIO genStakeKeyPair)
+  return $ BpiWallet kp skp k
 
 saveWallets :: MonadIO m => BpiWallet -> Maybe FilePath -> ReaderT ClusterEnv m (Either BpiError ())
 saveWallets bpiw fp = do
@@ -114,21 +123,27 @@ saveWallets bpiw fp = do
 
 -- | Save the wallet to a specific directory.
 saveWalletDir :: MonadIO m => BpiWallet -> FilePath -> m (Either BpiError ())
-saveWalletDir (BpiWallet pkh _ sk) wallDir = do
+saveWalletDir (BpiWallet pay stake _) wallDir = do
   liftIO $ createDirectoryIfMissing True wallDir
-  let pkhStr = Text.unpack (encodeByteString (fromBuiltin (LAPI.getPubKeyHash pkh)))
-      path = wallDir </> "signing-key-" ++ pkhStr <.> "skey"
-  res <- liftIO $ CAPI.writeFileTextEnvelope path (Just "Payment Signing Key") sk
-  return $ left (SignKeySaveError . show) res --todo: better error handling
+  pLogs <- liftIO $ writeKeyPair wallDir pay
+  sLogs <- maybe (pure []) (liftIO . writeStakeKeyPairs wallDir) stake
+
+  case listToMaybe $ dropWhile isRight (pLogs ++ sLogs) of
+    Nothing -> return $ pure ()
+    Just e -> return $ left (SignKeySaveError . show) e
 
 -- | Make `AnyAddress` for mainnet
 cardanoMainnetAddress :: BpiWallet -> AddressAny
-cardanoMainnetAddress (BpiWallet _ vk _) =
+cardanoMainnetAddress (BpiWallet pay stake _) =
   CAPI.toAddressAny $
-    CAPI.makeShelleyAddress
-      CAPI.Mainnet
-      (CAPI.PaymentCredentialByKey (CAPI.verificationKeyHash vk))
-      CAPI.NoStakeAddress
+    CAPI.ShelleyAddress
+      Shelley.Mainnet
+      ((\case (CAPI.PaymentKeyHash kh) -> Shelley.KeyHashObj kh) . CAPI.verificationKeyHash $ vKey pay)
+      ( maybe
+          Shelley.StakeRefNull
+          ((\case (CAPI.StakeKeyHash kh) -> Shelley.StakeRefBase $ Shelley.KeyHashObj kh) . CAPI.verificationKeyHash . sVKey)
+          stake
+      )
 
 -- | Get `String` representation of address on mainnet
 mkMainnetAddress :: BpiWallet -> String
@@ -137,5 +152,8 @@ mkMainnetAddress bw =
     . CAPI.serialiseAddress
     $ cardanoMainnetAddress bw
 
-ledgerPaymentPkh :: BpiWallet -> PaymentPubKeyHash
-ledgerPaymentPkh = PaymentPubKeyHash . walletPkh
+walletPaymentPkh :: BpiWallet -> PaymentPubKeyHash
+walletPaymentPkh = PaymentPubKeyHash . PubKeyHash . toBuiltin . CAPI.serialiseToRawBytes . CAPI.verificationKeyHash . vKey . payKeys
+
+walletStakePkh :: BpiWallet -> Maybe StakePubKeyHash
+walletStakePkh wall = StakePubKeyHash . PubKeyHash . toBuiltin . CAPI.serialiseToRawBytes . CAPI.verificationKeyHash . sVKey <$> stakeKeys wall
diff --git a/src/Test/Plutip/Internal/LocalCluster.hs b/src/Test/Plutip/Internal/LocalCluster.hs
index afb37846..ce1f8a83 100644
--- a/src/Test/Plutip/Internal/LocalCluster.hs
+++ b/src/Test/Plutip/Internal/LocalCluster.hs
@@ -274,7 +274,7 @@ launchChainIndex conf (RunningNode sp _block0 (netParams, _vData) _) dir = do
           & CIC.socketPath .~ nodeSocketFile sp
           & CIC.dbPath .~ dbPath
           & CIC.networkId .~ CAPI.Mainnet
-          & CIC.port .~ maybe (CIC.cicPort ChainIndex.defaultConfig) fromEnum (chainIndexPort conf)
+          & CIC.port .~ port
           & CIC.slotConfig .~ (def {scSlotLength = toMilliseconds slotLen})
 
   void $ async $ void $ ChainIndex.runMainWithLog (const $ return ()) config chainIndexConfig
diff --git a/src/Test/Plutip/LocalCluster.hs b/src/Test/Plutip/LocalCluster.hs
index eed89b8b..773b1415 100644
--- a/src/Test/Plutip/LocalCluster.hs
+++ b/src/Test/Plutip/LocalCluster.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+
 module Test.Plutip.LocalCluster (
   BpiWallet,
   addSomeWallet,
@@ -5,7 +7,7 @@ module Test.Plutip.LocalCluster (
   waitSeconds,
   mkMainnetAddress,
   cardanoMainnetAddress,
-  ledgerPaymentPkh,
+  walletPaymentPkh,
   withCluster,
   withConfiguredCluster,
   startCluster,
@@ -20,13 +22,16 @@ import Data.Default (def)
 import Data.List.NonEmpty (NonEmpty)
 import Numeric.Natural (Natural)
 import Test.Plutip.Config (PlutipConfig)
-import Test.Plutip.Contract (TestWallet (twInitDistribuition), TestWallets (unTestWallets), ada)
-import Test.Plutip.Internal.BotPlutusInterface.Wallet (
+import Test.Plutip.Contract (ClusterTest (ClusterTest), ada)
+import Test.Plutip.Internal.BotPlutusInterface.Types (
   BpiWallet,
+  TestWallet (TestWallet),
+ )
+import Test.Plutip.Internal.BotPlutusInterface.Wallet (
   addSomeWallet,
   cardanoMainnetAddress,
-  ledgerPaymentPkh,
   mkMainnetAddress,
+  walletPaymentPkh,
  )
 import Test.Plutip.Internal.LocalCluster (startCluster, stopCluster)
 import Test.Plutip.Internal.Types (ClusterEnv)
@@ -47,13 +52,13 @@ waitSeconds n = liftIO $ threadDelay (fromEnum n * 1_000_000)
 -- > test =
 -- >   withCluster
 -- >     "Tests with local cluster"
--- >     [ shouldSucceed "Get utxos" (initAda 100) $ const getUtxos
+-- >     [ assertExecution "Get utxos" (initAda (EntTag "w1") 100) (withContract $ const getUtxos) [shouldSucceed]]
 -- >     ...
 --
 -- @since 0.2
 withCluster ::
   String ->
-  [(TestWallets, IO (ClusterEnv, NonEmpty BpiWallet) -> TestTree)] ->
+  [ClusterTest] ->
   TestTree
 withCluster = withConfiguredCluster def
 
@@ -67,21 +72,21 @@ withCluster = withConfiguredCluster def
 -- >     let myConfig = PlutipConfig ...
 -- >     withConfiguredCluster myConfig
 -- >     "Tests with local cluster"
--- >     [ shouldSucceed "Get utxos" (initAda 100) $ const getUtxos
+-- >     [ assertExecution "Get utxos" (initAda (EntTag "w1") 100) (withContract $ const getUtxos) [shouldSucceed]]
 -- >     ...
 --
 -- @since 0.2
 withConfiguredCluster ::
   PlutipConfig ->
   String ->
-  [(TestWallets, IO (ClusterEnv, NonEmpty BpiWallet) -> TestTree)] ->
+  [ClusterTest] ->
   TestTree
 withConfiguredCluster conf name testCases =
   withResource (startCluster conf setup) (stopCluster . fst) $
     \getResource ->
       testGroup name $
         imap
-          (\idx (_, toTestGroup) -> toTestGroup $ second (!! idx) . snd <$> getResource)
+          (\idx (ClusterTest (_, toTestGroup)) -> toTestGroup $ second (!! idx) . snd <$> getResource)
           testCases
   where
     setup :: ReaderT ClusterEnv IO (ClusterEnv, [NonEmpty BpiWallet])
@@ -90,12 +95,17 @@ withConfiguredCluster conf name testCases =
 
       wallets <-
         traverse
-          (traverse addSomeWallet . fmap twInitDistribuition . unTestWallets . fst)
+          (traverse addTestWallet . getTestWallets)
           testCases
       -- had to bump waiting period here coz of chain-index slowdown,
       -- see https://github.com/mlabs-haskell/plutip/issues/120
       waitSeconds 5 -- wait for transactions to submit
       pure (env, wallets)
 
+    getTestWallets (ClusterTest (tws, _)) = tws
+
+    addTestWallet (TestWallet tag dist _) =
+      addSomeWallet tag dist
+
 imap :: (Int -> a -> b) -> [a] -> [b]
 imap fn = zipWith fn [0 ..]
diff --git a/src/Test/Plutip/Predicate.hs b/src/Test/Plutip/Predicate.hs
index 896e2bf9..7e11fc10 100644
--- a/src/Test/Plutip/Predicate.hs
+++ b/src/Test/Plutip/Predicate.hs
@@ -24,9 +24,9 @@ module Test.Plutip.Predicate (
 ) where
 
 import BotPlutusInterface.Types (TxBudget (TxBudget), mintBudgets, spendBudgets)
-import Data.List.NonEmpty (NonEmpty)
 import Data.Map (Map)
 import Data.Map qualified as Map
+import Data.Text (Text)
 import Ledger (ExBudget (ExBudget), ExCPU (ExCPU), ExMemory (ExMemory), TxId, Value)
 import PlutusCore.Evaluation.Machine.ExMemory (CostingInteger)
 import Prettyprinter (
@@ -62,10 +62,10 @@ data Predicate w e a = Predicate
     negative :: String
   , -- | some useful debugging info that `Predicate` can print based on contract execution result,
     -- used to print info in case of check failure
-    debugInfo :: ExecutionResult w e (a, NonEmpty Value) -> String
+    debugInfo :: ExecutionResult w e (a, Map Text Value) -> String
   , -- | check that `Predicate` performs on Contract execution result,
     -- if check evaluates to `False` test case considered failure
-    pCheck :: ExecutionResult w e (a, NonEmpty Value) -> Bool
+    pCheck :: ExecutionResult w e (a, Map Text Value) -> Bool
   }
 
 -- | `positive` description of `Predicate` that will be used as test case tag.
@@ -89,7 +89,7 @@ not predicate =
 -- | Check that Contract didn't fail.
 --
 -- @since 0.2
-shouldSucceed :: (Show e, Show a, Show w) => Predicate w e a
+shouldSucceed :: (Show w, Show e, Show a) => Predicate w e a
 shouldSucceed =
   Predicate
     "Contract should succeed"
@@ -98,7 +98,7 @@ shouldSucceed =
     isSuccessful
 
 -- | Pretty print ExecutionResult hiding budget stats and logs.
-prettyExecutionResult :: (Show e, Show w, Show a) => ExecutionResult e w a -> Doc ann
+prettyExecutionResult :: (Show e, Show w, Show a) => ExecutionResult w e a -> Doc ann
 prettyExecutionResult ExecutionResult {outcome, contractState} =
   vsep
     [ "Execution result {"
diff --git a/src/Test/Plutip/Tools/Address.hs b/src/Test/Plutip/Tools/Address.hs
index 3c3080d5..3145903b 100644
--- a/src/Test/Plutip/Tools/Address.hs
+++ b/src/Test/Plutip/Tools/Address.hs
@@ -50,10 +50,3 @@ ledgerToCardanoMainnet = Ledger.toCardanoAddressInEra CAPI.Mainnet
 ledgerToCardanoMainnet' :: Address.Address -> Either Ledger.ToCardanoError Text
 ledgerToCardanoMainnet' addr =
   CAPI.serialiseAddress <$> Ledger.toCardanoAddressInEra CAPI.Mainnet addr
-
--- | Get `String` representation of address on mainnet
--- mkMainnetAddress :: BpiWallet -> String
--- mkMainnetAddress bw =
---   unpack
---     . CAPI.serialiseAddress
---     $ cardanoMainnetAddress bw
diff --git a/test/Spec/Integration.hs b/test/Spec/Integration.hs
index c222311f..2b293b15 100644
--- a/test/Spec/Integration.hs
+++ b/test/Spec/Integration.hs
@@ -1,20 +1,25 @@
+{-# LANGUAGE OverloadedStrings #-}
+
 module Spec.Integration (test) where
 
 import BotPlutusInterface.Types (LogContext (ContractLog), LogLevel (Error), LogType (AnyLog))
 import Control.Exception (ErrorCall, Exception (fromException))
 import Control.Monad (void)
 import Data.Default (Default (def))
-import Data.List.NonEmpty (NonEmpty)
+import Data.List.NonEmpty (NonEmpty ((:|)))
 import Data.Map qualified as Map
 import Data.Maybe (isJust)
 import Data.Text (Text, isInfixOf, pack)
+import Ledger (Address (Address), PaymentPubKeyHash (PaymentPubKeyHash), StakePubKeyHash (StakePubKeyHash))
 import Ledger.Ada (lovelaceValueOf)
 import Ledger.Constraints (MkTxError (OwnPubKeyMissing))
 import Plutus.Contract (
   ContractError (ConstraintResolutionContractError),
+  throwError,
   waitNSlots,
  )
 import Plutus.Contract qualified as Contract
+import Plutus.V2.Ledger.Api (Credential (PubKeyCredential), StakingCredential (StakingHash))
 import Spec.TestContract.AdjustTx (runAdjustTest)
 import Spec.TestContract.AlwaysFail (lockThenFailToSpend)
 import Spec.TestContract.LockSpendMint (lockThenSpend)
@@ -26,11 +31,11 @@ import Spec.TestContract.SimpleContracts (
   ownValue,
   ownValueToState,
   payTo,
+  payToPubKeyAddress,
  )
 import Spec.TestContract.ValidateTimeRange (failingTimeContract, successTimeContract)
 import Test.Plutip.Contract (
-  TestWallets,
-  ValueOrdering (VLt),
+  ClusterTest,
   assertExecution,
   assertExecutionWith,
   initAda,
@@ -42,12 +47,14 @@ import Test.Plutip.Contract (
   withContract,
   withContractAs,
  )
+import Test.Plutip.Contract.Types (WalletTag (BaseTag, EntTag))
+import Test.Plutip.Internal.BotPlutusInterface.Lookups (WalletLookups (lookupWallet), lookupAddress)
+import Test.Plutip.Internal.BotPlutusInterface.Types (BaseWallet (BaseWallet), EntWallet (EntWallet), ValueOrdering (VLt))
 import Test.Plutip.Internal.Types (
-  ClusterEnv,
   FailureReason (CaughtException, ContractExecutionError),
   isException,
  )
-import Test.Plutip.LocalCluster (BpiWallet, withConfiguredCluster)
+import Test.Plutip.LocalCluster (withConfiguredCluster)
 import Test.Plutip.Options (TraceOption (ShowBudgets, ShowTraceButOnlyContext))
 import Test.Plutip.Predicate (
   assertOverallBudget,
@@ -77,14 +84,14 @@ test =
         -- Basic Succeed or Failed tests
         assertExecution
           "Contract 1"
-          (initAda (100 : replicate 10 7))
+          (initAda (EntTag "w1") (100 : replicate 10 7))
           (withContract $ const getUtxos)
           [ shouldSucceed
           , Predicate.not shouldFail
           ]
       , assertExecution
           "Contract 2"
-          (initAda [100])
+          (initAda (EntTag "w1") [100])
           (withContract $ const getUtxosThrowsErr)
           [ shouldFail
           , Predicate.not shouldSucceed
@@ -92,7 +99,7 @@ test =
       , assertExecutionWith
           [ShowTraceButOnlyContext ContractLog $ Error [AnyLog]]
           "Contract 3"
-          (initAda [100])
+          (initAda (EntTag "w1") [100])
           ( withContract $
               const $ do
                 Contract.logInfo @Text "Some contract log with Info level."
@@ -103,39 +110,50 @@ test =
           ]
       , assertExecution
           "Pay negative amount"
-          (initAda [100])
-          (withContract $ \[pkh1] -> payTo pkh1 (-10_000_000))
+          (initAda (EntTag "w1") [100])
+          ( withContract $ \ws -> do
+              EntWallet pkh1 <- lookupWallet ws (EntTag "w1")
+              payTo pkh1 (-10_000_000)
+          )
           [shouldFail]
       , -- Tests with wallet's Value assertions
         assertExecution
           "Pay from wallet to wallet"
-          (initAda [100] <> initAndAssertAda [100, 13] 123)
-          (withContract $ \[pkh1] -> payTo pkh1 10_000_000)
+          ( initAda (EntTag "w1") [100]
+              <> initAndAssertAda (EntTag "w2") [100, 13] 123
+          )
+          ( withContract $ \ws -> do
+              EntWallet pkh1 <- lookupWallet ws (EntTag "w2")
+              payTo pkh1 10_000_000
+          )
           [shouldSucceed]
       , assertExecution
           "Two contracts one after another"
-          ( initAndAssertAdaWith [100] VLt 100 -- own wallet (index 0 in wallets list)
-              <> initAndAssertAdaWith [100] VLt 100 -- wallet with index 1 in wallets list
+          ( initAndAssertAdaWith (EntTag "w0") [100] VLt 100 -- own wallet (index 0 in wallets lookups)
+              <> initAndAssertAdaWith (EntTag "w1") [100] VLt 100 -- wallet with index 1 in wallets lookups
           )
           ( do
               void $ -- run something prior to the contract which result will be checked
-                withContract $
-                  \[pkh1] -> payTo pkh1 10_000_000
-              withContractAs 1 $ -- run contract which result will be checked
-                \[pkh1] -> payTo pkh1 10_000_000
+                withContract $ \ws -> do
+                  addr1 <- lookupAddress ws "w1"
+                  payToPubKeyAddress addr1 10_000_000
+              withContractAs "w1" $ -- run contract which result will be checked
+                \ws -> do
+                  addr0 <- lookupAddress ws "w0"
+                  payToPubKeyAddress addr0 10_000_000
           )
           [shouldSucceed]
       , -- Tests with assertions on Contract return value
         assertExecution
           "Initiate wallet and get UTxOs"
-          (initAda [100])
+          (initAda (BaseTag "") [100])
           (withContract $ const getUtxos)
           [ yieldSatisfies "Returns single UTxO" ((== 1) . Map.size)
           ]
       , let initFunds = 10_000_000
          in assertExecution
               "Should yield own initial Ada"
-              (initLovelace [toEnum initFunds])
+              (initLovelace (BaseTag "") [toEnum initFunds])
               (withContract $ const ownValue)
               [ shouldYield (lovelaceValueOf $ toEnum initFunds)
               ]
@@ -143,7 +161,7 @@ test =
         let initFunds = 10_000_000
          in assertExecution
               "Puts own UTxOs Value to state"
-              (initLovelace [toEnum initFunds])
+              (initLovelace (BaseTag "") [toEnum initFunds])
               (withContract $ const ownValueToState)
               [ stateIs [lovelaceValueOf $ toEnum initFunds]
               , Predicate.not $ stateSatisfies "length > 1" ((> 1) . length)
@@ -155,7 +173,7 @@ test =
               _ -> False
          in assertExecution
               ("Contract which throws `" <> show expectedErr <> "`")
-              (initAda [100])
+              (initAda (BaseTag "") [100])
               (withContract $ const getUtxosThrowsErr)
               [ shouldThrow expectedErr
               , errorSatisfies "Throws resolution error" isResolutionError
@@ -166,7 +184,7 @@ test =
               _ -> False
          in assertExecution
               "Contract which throws exception"
-              (initAda [100])
+              (initAda (EntTag "") [100])
               (withContract $ const getUtxosThrowsEx)
               [ shouldFail
               , Predicate.not shouldSucceed
@@ -176,7 +194,7 @@ test =
         assertExecutionWith
           [ShowBudgets] -- this influences displaying the budgets only and is not necessary for budget assertions
           "Lock then spend contract"
-          (initAda (replicate 3 300))
+          (initAda (EntTag "") (replicate 3 300))
           (withContract $ const lockThenSpend)
           [ shouldSucceed
           , budgetsFitUnder
@@ -194,25 +212,26 @@ test =
               _ -> False
          in assertExecution
               "Fails because outside validity interval"
-              (initAda [100])
+              (initAda (EntTag "") [100])
               (withContract $ const failingTimeContract)
               [ shouldFail
               , failReasonSatisfies "Execution error is OutsideValidityIntervalUTxO" isValidityError
               ]
       , assertExecution
           "Passes validation with exact time range checks"
-          (initAda [100])
+          (initAda (EntTag "") [100])
           (withContract $ const successTimeContract)
           [shouldSucceed]
       , -- always fail validation test
         let errCheck e = "I always fail" `isInfixOf` pack (show e)
          in assertExecution
               "Always fails to validate"
-              (initAda [100])
+              (initAda (EntTag "") [100])
               (withContract $ const lockThenFailToSpend)
               [ shouldFail
               , errorSatisfies "Fail validation with 'I always fail'" errCheck
               ]
+      , walletLookupsTest
       , -- Test `adjustUnbalancedTx`
         runAdjustTest
       , testBugMintAndPay
@@ -220,18 +239,20 @@ test =
       ++ testValueAssertionsOrderCorrectness
 
 -- https://github.com/mlabs-haskell/plutip/issues/138
-testBugMintAndPay :: (TestWallets, IO (ClusterEnv, NonEmpty BpiWallet) -> TestTree)
+testBugMintAndPay ::ClusterTest
 testBugMintAndPay =
   assertExecution
     "Adjustment of outputs with 0 Ada does not fail"
-    (withCollateral $ initAda [1000] <> initAda [1111])
-    (withContract $ \[p1] -> zeroAdaOutTestContract p1)
+    (withCollateral $ initAda (EntTag "w0") [1000] <> initAda (EntTag "w1") [1111])
+    (withContract $ \ws -> do
+      EntWallet w1pkh <- lookupWallet ws (EntTag "w1")
+      zeroAdaOutTestContract w1pkh
+    )
     [ shouldSucceed
     ]
 
 -- Tests for https://github.com/mlabs-haskell/plutip/issues/84
-testValueAssertionsOrderCorrectness ::
-  [(TestWallets, IO (ClusterEnv, NonEmpty BpiWallet) -> TestTree)]
+testValueAssertionsOrderCorrectness :: [ClusterTest]
 testValueAssertionsOrderCorrectness =
   [ -- withContract case
     let wallet0 = 100_000_000
@@ -252,12 +273,14 @@ testValueAssertionsOrderCorrectness =
      in assertExecution
           "Values asserted in correct order with withContract"
           ( withCollateral $
-              initAndAssertLovelace [wallet0] wallet0After
-                <> initAndAssertLovelace [wallet1] wallet1After
-                <> initAndAssertLovelace [wallet2] wallet2After
+              initAndAssertLovelace (EntTag "w0") [wallet0] wallet0After
+                <> initAndAssertLovelace (EntTag "w1") [wallet1] wallet1After
+                <> initAndAssertLovelace (EntTag "w2") [wallet2] wallet2After
           )
           ( do
-              withContract $ \[w1pkh, w2pkh] -> do
+              withContract $ \ws -> do
+                EntWallet w1pkh <- lookupWallet ws (EntTag "w1")
+                EntWallet w2pkh <- lookupWallet ws (EntTag "w2")
                 _ <- payTo w1pkh (toInteger payTo1Amt)
                 _ <- waitNSlots 2
                 payTo w2pkh (toInteger payTo2Amt)
@@ -290,19 +313,57 @@ testValueAssertionsOrderCorrectness =
      in assertExecution
           "Values asserted in correct order with withContractAs"
           ( withCollateral $ -- Initialize all the wallets with the collateral utxo.
-              initAndAssertLovelace [wallet0] wallet0After
-                <> initAndAssertLovelace [wallet1] wallet1After
-                <> initAndAssertLovelace [wallet2] wallet2After
+              initAndAssertLovelace (EntTag "w0") [wallet0] wallet0After
+                <> initAndAssertLovelace (EntTag "w1") [wallet1] wallet1After
+                <> initAndAssertLovelace (EntTag "w2") [wallet2] wallet2After
           )
           ( do
               void $
-                withContractAs 1 $ \[w0pkh, w2pkh] -> do
+                withContractAs "w1" $ \ws -> do
+                  EntWallet w0pkh <- lookupWallet ws (EntTag "w0")
+                  EntWallet w2pkh <- lookupWallet ws (EntTag "w2")
                   _ <- payTo w0pkh (toInteger payTo0Amt)
                   _ <- waitNSlots 2
                   payTo w2pkh (toInteger payTo2Amt)
 
-              withContractAs 2 $ \[_, w1pkh] -> do
+              withContractAs "w2" $ \ws -> do
+                EntWallet w1pkh <- lookupWallet ws (EntTag "w1")
                 payTo w1pkh (toInteger payTo1Amt)
           )
           [shouldSucceed]
   ]
+
+walletLookupsTest :: ClusterTest
+walletLookupsTest =
+  assertExecution @() @Text
+    "Wallets initilized expectedly."
+    ( initAndAssertAda (BaseTag "a") [10, 20] 30
+        <> initAndAssertAda (BaseTag "b") [11, 22] 33
+        <> initAndAssertAda (EntTag "c") [1, 2] 3
+    )
+    ( withContract $ \ws -> do
+        BaseWallet pkhb spkhb <- lookupWallet ws (BaseTag "b")
+        EntWallet pkhc <- lookupWallet ws (EntTag "c")
+        addrb <- lookupAddress ws "b"
+        addrc <- lookupAddress ws "c"
+
+        case addrb of
+          Address (PubKeyCredential pkhb') (Just (StakingHash (PubKeyCredential spkhb'))) ->
+            if pkhb == PaymentPubKeyHash pkhb' && spkhb == StakePubKeyHash spkhb'
+              then pure ()
+              else throwError "Unexpected key hashes of wallet 'b'."
+          _ -> throwError "Unexpected address of wallet 'b'."
+
+        case addrc of
+          Address (PubKeyCredential pkhc') Nothing ->
+            if pkhc == PaymentPubKeyHash pkhc'
+              then pure ()
+              else throwError "Unexpected key hashes of wallet 'c'."
+          _ -> throwError "Unexpected address of wallet 'c'."
+
+        ourAddr :| _ <- Contract.ownAddresses
+        case ourAddr of
+          Address (PubKeyCredential _) (Just (StakingHash (PubKeyCredential _))) -> pure ()
+          _ -> throwError "Unexpected contract own address."
+    )
+    [shouldSucceed]
diff --git a/test/Spec/TestContract/AdjustTx.hs b/test/Spec/TestContract/AdjustTx.hs
index 54633b3f..6ea6edf1 100644
--- a/test/Spec/TestContract/AdjustTx.hs
+++ b/test/Spec/TestContract/AdjustTx.hs
@@ -3,7 +3,6 @@ module Spec.TestContract.AdjustTx (
 ) where
 
 import Control.Lens.Operators ((^.))
-import Data.List.NonEmpty qualified as NonEmpty
 import Data.Text (Text)
 import Data.Void (Void)
 import Ledger (
@@ -26,18 +25,17 @@ import Plutus.Contract (
 import Plutus.Contract qualified as Contract
 import Plutus.PAB.Effects.Contract.Builtin (EmptySchema)
 import Test.Plutip.Contract (
-  TestWallets,
+  ClusterTest,
   assertExecution,
   initAda,
   withContract,
  )
-import Test.Plutip.Internal.BotPlutusInterface.Wallet (BpiWallet)
-import Test.Plutip.Internal.Types (ClusterEnv)
+import Test.Plutip.Internal.BotPlutusInterface.Lookups (WalletLookups (lookupWallet))
+import Test.Plutip.Internal.BotPlutusInterface.Types (EntWallet (EntWallet), WalletTag (EntTag))
 import Test.Plutip.Predicate (
   shouldSucceed,
   yieldSatisfies,
  )
-import Test.Tasty (TestTree)
 import Prelude
 
 adjustTx :: PaymentPubKeyHash -> Contract () EmptySchema Text [Value]
@@ -66,17 +64,17 @@ adjustTx' [] = do
   adjustTx pkh
 adjustTx' (pkh : _) = adjustTx pkh
 
--- | A type for the output of `assertExecution`.
-type PlutipTest = (TestWallets, IO (ClusterEnv, NonEmpty.NonEmpty BpiWallet) -> TestTree)
-
 -- | Tests whether `adjustUnbalancedTx` actually tops up the
 -- UTxO to get to the minimum required ADA.
-runAdjustTest :: PlutipTest
+runAdjustTest :: ClusterTest
 runAdjustTest =
   assertExecution
     "Adjust Unbalanced Tx Contract"
-    (initAda [1000] <> initAda [1000])
-    (withContract adjustTx')
+    (initAda (EntTag "w1") [1000] <> initAda (EntTag "w2") [1000])
+    ( withContract $ \ws -> do
+        EntWallet pkh <- lookupWallet ws (EntTag "w2")
+        adjustTx' [pkh]
+    )
     [ shouldSucceed
     , yieldSatisfies
         "All UTxOs have minimum(?) ADA."
diff --git a/test/Spec/TestContract/SimpleContracts.hs b/test/Spec/TestContract/SimpleContracts.hs
index 5f9d57aa..36386622 100644
--- a/test/Spec/TestContract/SimpleContracts.hs
+++ b/test/Spec/TestContract/SimpleContracts.hs
@@ -5,20 +5,24 @@ module Spec.TestContract.SimpleContracts (
   payTo,
   ownValue,
   ownValueToState,
+  payToPubKeyAddress,
 ) where
 
 import Plutus.Contract (
   Contract,
   ContractError (ConstraintResolutionContractError),
   submitTx,
+  throwError,
   utxosAt,
  )
 import Plutus.Contract qualified as Contract
 
 import Ledger (
+  Address (Address),
   CardanoTx,
   ChainIndexTxOut,
-  PaymentPubKeyHash,
+  PaymentPubKeyHash (PaymentPubKeyHash),
+  StakePubKeyHash (StakePubKeyHash),
   TxOutRef,
   Value,
   ciTxOutValue,
@@ -35,6 +39,8 @@ import Data.Text (Text)
 import Ledger.Constraints (MkTxError (OwnPubKeyMissing))
 import Ledger.Constraints qualified as Constraints
 import Plutus.PAB.Effects.Contract.Builtin (EmptySchema)
+import Plutus.V1.Ledger.Api (Credential (ScriptCredential), StakingCredential (StakingHash, StakingPtr))
+import Plutus.V1.Ledger.Credential (Credential (PubKeyCredential))
 
 getUtxos :: Contract [Value] EmptySchema Text (Map TxOutRef ChainIndexTxOut)
 getUtxos = do
@@ -54,6 +60,27 @@ payTo toPkh amt = do
   _ <- Contract.awaitTxConfirmed (getCardanoTxId tx)
   pure tx
 
+payToPubKeyAddress :: Address -> Integer -> Contract () EmptySchema Text CardanoTx
+payToPubKeyAddress (Address crd stake) amt = do
+  pkh <- case crd of
+    PubKeyCredential pkh -> pure pkh
+    ScriptCredential _ -> throwError "Expected PubKey credential."
+  mspkh <- case stake of
+    Just (StakingHash (PubKeyCredential spkh)) -> pure $ Just spkh
+    Just (StakingHash (ScriptCredential _)) -> throwError "Expected PubKey credential."
+    Just StakingPtr {} -> throwError "No support for staking pointers."
+    Nothing -> pure Nothing
+
+  let constr =
+        maybe
+          (Constraints.mustPayToPubKey (PaymentPubKeyHash pkh) (Ada.lovelaceValueOf amt))
+          (\spkh -> Constraints.mustPayToPubKeyAddress (PaymentPubKeyHash pkh) (StakePubKeyHash spkh) (Ada.lovelaceValueOf amt))
+          mspkh
+
+  tx <- submitTx constr
+  _ <- Contract.awaitTxConfirmed (getCardanoTxId tx)
+  pure tx
+
 ownValue :: Contract [Value] EmptySchema Text Value
 ownValue = foldMap (^. ciTxOutValue) <$> getUtxos