Testing if one monad is called in the wrong environment

I have a MonadReader that looks like the following:

data Environment = EnvironmentA | EnvironmentB

mainMonad ::
  ( MonadReader m
  )
    => m Type
mainMonad = do
  env <- ask
  case env of
    EnvironmentA -> monadA
    EnvironmentB -> monadB

monadA ::
  ( MonadReader m
  )
    => m Type
monadA = do
  ...

monadB ::
  ( MonadReader m
  )
    => m Type
monadB = do
  ...

Basically it our main monad uses the environment to switch between monadA and monadB. Of course my actual code has a more complex Environment type and more than just two monads that can be called.

Now between monadA and monadB there may be a number of calls to each other. But each time we use local to adjust the monad to the proper environment so that monadA.

monadA ::
  ( MonadReader m
  )
    => m Type
monadA = do
  ...
  bResult <- local (const EnvironmentB) monadB
  ...

monadB ::
  ( MonadReader m
  )
    => m Type
monadB = do
  ...
  aResult <- local (const EnvironmentA) monadA
  ...

Now my issue is that it can be pretty easy to miss a local in my code and just call a monad directly with the wrong environment. The results of doing this are often rather subtle and rather varied. This ends up making the symptoms of the problem rather hard to catch with unit testing. So I would like to target the problem directly by adding a clause to my unit test that says something along the lines of:

Call mainMonad check that over the course of evaluating it we never have a monad called with the wrong environment.

That way I can catch these mistakes without having to comb through the code very carefully. Now after thinking about this for a little while I have not come up with a very neat way to do this. I've thought of a couple of ways that do work but I am not quite happy with:

• Hard crash when called with the wrong environment. I could fix this by adding a condition to the front of each monad that hard crashes if it detects it being called with the wrong environment. For example:

  monadA ::
   ( MonadReader m
   )
     => m Type
  monadA = do
    env <- ask
    case env of
      EnvironmentA -> return ()
      _ -> undefined
    ...
  
  

  This is not ideal since I would really prefer the customer to experience the slight issues caused by calling things with the wrong environment rather than a hard crash in the event that the test handler does not catch the issue. It sort of seems like the nuclear option.

• Fiddling with types. I also tried changing the types of monadA and monadB so that monadA could not be called directly from monadB or vice versa. This is very nice in that it catches the problems at compile time. This has the issue of being a bit of a pain to maintain, and it is quite complex. Since monadA and monadB may each share a number of common monads of the type (MonadReader m) => m Type each and every one of those has to be lifted as well. Really it pretty much guarantees that every line now has a lift. I'm not opposed to type based solutions but I don't want to have to spend a huge deal of time just maintaining a unit test.

So is there a way I can ensure that my monads are always called from the correct environment?