Repository: commanded/commanded (elixir)

Issue:
Replace aggregate lifespan `after_command/1` callback with `after_event/1`
An aggregate process lifespan should be controlled by the domain events it creates, not the commands, since a command may be rejected or ignored.

The events created by an aggregate are returned when executing a command so it should be straightforward to call an `after_event/1` callback and pass it the last domain event.

Pull request:
Replace aggregate lifespan `after_command/1` callback with `after_event/1`
An aggregate process lifespan should be controlled by the domain events it creates, not the commands, since a command may be rejected or ignored.

The events created by an aggregate are returned when executing a command so it should be straightforward to call an after_event/1 callback and pass it the last domain event.

#### Supported return values

- Non-negative integer - specify an inactivity timeout, in millisconds.
- `:infinity` - prevent the aggregate instance from shutting down.
- `:hibernate` - send the process into hibernation.
- `:stop` - immediately shutdown the aggregate process.

Closes #126.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2564b2a..8ed6b52 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,11 @@
 - Event handler `error/3` callback ([#133](https://github.com/commanded/commanded/pull/133)).
 - Support distributed dispatch consistency ([#135](https://github.com/commanded/commanded/pull/135)).
 - Defer event handler and process router init until after subscribed ([#138](https://github.com/commanded/commanded/pull/138)).
+- Replace aggregate lifespan `after_command/1` callback with `after_event/1` ([#139](https://github.com/commanded/commanded/issues/139)).
+
+### Breaking changes
+
+- The `Commanded.Aggregates.AggregateLifespan` behaviour has been changed from `after_command/1` to `after_event/1`. You will need to update your own lifespan modules to use events instead of commands to shutdown an aggregate process after an inactivity timeout.
 
 ## v0.15.1
 
diff --git a/guides/Commands.md b/guides/Commands.md
index 6563d81..fc85c58 100644
--- a/guides/Commands.md
+++ b/guides/Commands.md
@@ -268,12 +268,15 @@ You should always use string keys in your metadata map; atom keys will be conver
 
 By default an aggregate instance process will run indefinitely once started. You can control this by implementing the `Commanded.Aggregates.AggregateLifespan` behaviour in a module.
 
+Define a module that implements the `Commanded.Aggregates.AggregateLifespan` behaviour:
+
 ```elixir
 defmodule BankAccountLifespan do
   @behaviour Commanded.Aggregates.AggregateLifespan
 
-  def after_command(%OpenAccount{}), do: :infinity
-  def after_command(%CloseAccount{}), do: 0
+  def after_event(%BankAccountOpened{}), do: :infinity
+  def after_event(%MoneyDeposited{}), do: 60_000
+  def after_event(%BankAccountClosed{}), do: :stop
 end
 ```
 
@@ -290,11 +293,11 @@ defmodule BankRouter do
 end
 ```
 
-The timeout is specified in milliseconds, after which time the aggregate process will be stopped if no other messages are received.
+The inactivity timeout is specified in milliseconds, after which time the aggregate process will be stopped if no other messages are received.
 
-You can also return `:hibernate` and the process is hibernated, it will continue its loop once a message is in its message queue. Hibernating an aggregate causes garbage collection and minimises the memory used by the process. Hibernating should not be used aggressively as too much time could be spent garbage collecting.
+Return `:stop` to immediately shutdown the aggregate process. Return `:infinity` to prevent the aggregate instance from shutting down.
 
-Return `:infinity` to keep the aggregate instance process running indefinitely.
+You can also return `:hibernate` and the process is hibernated, it will continue its loop once a message is in its message queue. Hibernating an aggregate causes garbage collection and minimises the memory used by the process. Hibernating should not be used aggressively as too much time could be spent garbage collecting.
 
 ## Middleware
 
diff --git a/lib/commanded/aggregates/aggregate.ex b/lib/commanded/aggregates/aggregate.ex
index e81dfa0..49928de 100644
--- a/lib/commanded/aggregates/aggregate.ex
+++ b/lib/commanded/aggregates/aggregate.ex
@@ -42,22 +42,20 @@ defmodule Commanded.Aggregates.Aggregate do
 
   require Logger
 
-  alias Commanded.Aggregates.{Aggregate,ExecutionContext}
+  alias Commanded.Aggregates.{Aggregate, ExecutionContext}
   alias Commanded.Event.Mapper
   alias Commanded.EventStore
   alias Commanded.EventStore.SnapshotData
 
   @read_event_batch_size 100
 
-  defstruct [
-    aggregate_module: nil,
-    aggregate_uuid: nil,
-    aggregate_state: nil,
-    aggregate_version: 0,
-    snapshot_every: nil,
-    snapshot_module_version: 1,
-    snapshot_version: 0,
-  ]
+  defstruct aggregate_module: nil,
+            aggregate_uuid: nil,
+            aggregate_state: nil,
+            aggregate_version: 0,
+            snapshot_every: nil,
+            snapshot_module_version: 1,
+            snapshot_version: 0
 
   def start_link(aggregate_module, aggregate_uuid, opts \\ []) do
     snapshot_options = snapshot_options(aggregate_module)
@@ -66,16 +64,14 @@ defmodule Commanded.Aggregates.Aggregate do
       aggregate_module: aggregate_module,
       aggregate_uuid: aggregate_uuid,
       snapshot_every: Keyword.get(snapshot_options, :snapshot_every),
-      snapshot_module_version: Keyword.get(snapshot_options, :snapshot_version, 1),
+      snapshot_module_version: Keyword.get(snapshot_options, :snapshot_version, 1)
     }
 
     GenServer.start_link(__MODULE__, aggregate, opts)
   end
 
-
   @doc false
-  def name(aggregate_module, aggregate_uuid),
-    do: {aggregate_module, aggregate_uuid}
+  def name(aggregate_module, aggregate_uuid), do: {aggregate_module, aggregate_uuid}
 
   def init(%Aggregate{} = state) do
     # initial aggregate state is populated by loading its state snapshot and/or
@@ -107,11 +103,16 @@ defmodule Commanded.Aggregates.Aggregate do
 
   """
   def execute(aggregate_module, aggregate_uuid, %ExecutionContext{} = context, timeout \\ 5_000) do
-    GenServer.call(via_name(aggregate_module, aggregate_uuid), {:execute_command, context}, timeout)
+    GenServer.call(
+      via_name(aggregate_module, aggregate_uuid),
+      {:execute_command, context},
+      timeout
+    )
   end
 
   @doc false
   def aggregate_state(aggregate_module, aggregate_uuid, timeout \\ 5_000)
+
   def aggregate_state(aggregate_module, aggregate_uuid, timeout) do
     GenServer.call(via_name(aggregate_module, aggregate_uuid), :aggregate_state, timeout)
   end
@@ -140,14 +141,27 @@ defmodule Commanded.Aggregates.Aggregate do
   def handle_call({:execute_command, %ExecutionContext{} = context}, _from, %Aggregate{} = state) do
     {reply, state} = execute_command(context, state)
 
-    lifespan_timeout = aggregate_lifespan_timeout(context)
+    lifespan_timeout =
+      case reply do
+        {:ok, _stream_version, events} ->
+          aggregate_lifespan_timeout(context, events)
 
-    if snapshotting_enabled?(state) && snapshot_required?(state) do
-      :ok = GenServer.cast(self(), {:snapshot_state, lifespan_timeout})
+        _ ->
+          :infinity
+      end
 
-      {:reply, reply, state}
-    else
-      {:reply, reply, state, lifespan_timeout}
+    case lifespan_timeout do
+      :stop ->
+        {:stop, :normal, reply, state}
+
+      _ ->
+        if snapshotting_enabled?(state) && snapshot_required?(state) do
+          :ok = GenServer.cast(self(), {:snapshot_state, lifespan_timeout})
+
+          {:reply, reply, state}
+        else
+          {:reply, reply, state, lifespan_timeout}
+        end
     end
   end
 
@@ -185,17 +199,15 @@ defmodule Commanded.Aggregates.Aggregate do
            {:ok, snapshot} <- EventStore.read_snapshot(aggregate_uuid),
            true <- snapshot_valid?(snapshot, state) do
         # populate initial state from snapshot
-        %Aggregate{state |
-          aggregate_version: snapshot.source_version,
-          aggregate_state: snapshot.data
+        %Aggregate{
+          state
+          | aggregate_version: snapshot.source_version,
+            aggregate_state: snapshot.data
         }
       else
         _ ->
           # no snapshot (or outdated), use intial empty state
-          %Aggregate{state |
-            aggregate_version: 0,
-            aggregate_state: struct(aggregate_module)
-          }
+          %Aggregate{state | aggregate_version: 0, aggregate_state: struct(aggregate_module)}
       end
 
     rebuild_from_events(aggregate)
@@ -249,38 +261,51 @@ defmodule Commanded.Aggregates.Aggregate do
     end
   end
 
-  defp aggregate_lifespan_timeout(%ExecutionContext{command: command, lifespan: lifespan}) do
-    case lifespan.after_command(command) do
-      :infinity ->
-        :infinity
+  defp aggregate_lifespan_timeout(_context, []), do: :infinity
 
-      :hibernate ->
-        :hibernate
+  defp aggregate_lifespan_timeout(%ExecutionContext{} = context, [event]) do
+    %ExecutionContext{lifespan: lifespan} = context
+
+    case lifespan.after_event(event) do
+      timeout when timeout in [:infinity, :hibernate, :stop] ->
+        timeout
 
       timeout when is_integer(timeout) and timeout >= 0 ->
         timeout
 
       invalid ->
         Logger.warn(fn ->
-          "Invalid timeout for aggregate lifespan #{inspect lifespan}, expected a non-negative integer, `:infinity`, or `:hibernate` but got: #{inspect invalid}"
+          "Invalid timeout for aggregate lifespan #{inspect(lifespan)}, expected a non-negative integer, `:infinity`, or `:hibernate` but got: #{
+            inspect(invalid)
+          }"
         end)
 
         :infinity
     end
   end
 
+  defp aggregate_lifespan_timeout(context, [_event | events]),
+    do: aggregate_lifespan_timeout(context, events)
+
   # is snapshotting configured for the aggregate?
   defp snapshotting_enabled?(%Aggregate{snapshot_every: snapshot_every}),
     do: is_number(snapshot_every) && snapshot_every > 0
 
   # was the snapshot taken at the current version?
-  defp snapshot_valid?(%SnapshotData{metadata: metadata}, %Aggregate{snapshot_module_version: expected_version}) do
+  defp snapshot_valid?(%SnapshotData{metadata: metadata}, %Aggregate{
+         snapshot_module_version: expected_version
+       }) do
     Map.get(metadata, "snapshot_version", 1) == expected_version
   end
 
   # take a snapshot now?
-  defp snapshot_required?(%Aggregate{aggregate_version: aggregate_version, snapshot_every: snapshot_every, snapshot_version: snapshot_version})
-    when aggregate_version - snapshot_version >= snapshot_every, do: true
+  defp snapshot_required?(%Aggregate{
+         aggregate_version: aggregate_version,
+         snapshot_every: snapshot_every,
+         snapshot_version: snapshot_version
+       })
+       when aggregate_version - snapshot_version >= snapshot_every,
+       do: true
 
   # not yet enough events to snapshot
   defp snapshot_required?(_state), do: false
@@ -291,7 +316,7 @@ defmodule Commanded.Aggregates.Aggregate do
       aggregate_module: aggregate_module,
       aggregate_uuid: aggregate_uuid,
       aggregate_version: aggregate_version,
-      aggregate_state: aggregate_state,
+      aggregate_state: aggregate_state
     } = state
 
     Logger.debug(fn -> describe(state) <> " recording snapshot" end)
@@ -310,7 +335,8 @@ defmodule Commanded.Aggregates.Aggregate do
   end
 
   defp execute_command(%ExecutionContext{retry_attempts: retry_attempts}, %Aggregate{} = state)
-    when retry_attempts < 0, do: {{:error, :too_many_attempts}, state}
+       when retry_attempts < 0,
+       do: {{:error, :too_many_attempts}, state}
 
   defp execute_command(%ExecutionContext{} = context, %Aggregate{} = state) do
     %ExecutionContext{
@@ -326,7 +352,7 @@ defmodule Commanded.Aggregates.Aggregate do
       aggregate_state: aggregate_state
     } = state
 
-    Logger.debug(fn -> describe(state) <> " executing command: #{inspect command}" end)
+    Logger.debug(fn -> describe(state) <> " executing command: #{inspect(command)}" end)
 
     {reply, state} =
       case Kernel.apply(handler, function, [aggregate_state, command]) do
@@ -374,10 +400,12 @@ defmodule Commanded.Aggregates.Aggregate do
   defp persist_events(pending_events, aggregate_state, context, %Aggregate{} = state) do
     %Aggregate{aggregate_uuid: aggregate_uuid, aggregate_version: expected_version} = state
 
-    with {:ok, stream_version} <- append_to_stream(pending_events, aggregate_uuid, expected_version, context) do
-      state = %Aggregate{state |
-        aggregate_state: aggregate_state,
-        aggregate_version: stream_version,
+    with {:ok, stream_version} <-
+           append_to_stream(pending_events, aggregate_uuid, expected_version, context) do
+      state = %Aggregate{
+        state
+        | aggregate_state: aggregate_state,
+          aggregate_version: stream_version
       }
 
       {{:ok, stream_version, pending_events}, state}
@@ -388,6 +416,7 @@ defmodule Commanded.Aggregates.Aggregate do
   end
 
   defp append_to_stream([], _stream_uuid, expected_version, _context), do: {:ok, expected_version}
+
   defp append_to_stream(pending_events, stream_uuid, expected_version, context) do
     %ExecutionContext{
       causation_id: causation_id,
@@ -414,6 +443,6 @@ defmodule Commanded.Aggregates.Aggregate do
       aggregate_version: aggregate_version
     } = aggregate
 
-    "#{inspect aggregate_module}<#{aggregate_uuid}@#{aggregate_version}>"
+    "#{inspect(aggregate_module)}<#{aggregate_uuid}@#{aggregate_version}>"
   end
 end
diff --git a/lib/commanded/aggregates/aggregate_lifespan.ex b/lib/commanded/aggregates/aggregate_lifespan.ex
index b7bf8f1..4e3453d 100644
--- a/lib/commanded/aggregates/aggregate_lifespan.ex
+++ b/lib/commanded/aggregates/aggregate_lifespan.ex
@@ -1,23 +1,62 @@
 defmodule Commanded.Aggregates.AggregateLifespan do
   @moduledoc """
-  The `Commanded.Aggregates.AggregateLifespan` behaviour is used to control an aggregate lifespan.
+  The `Commanded.Aggregates.AggregateLifespan` behaviour is used to control an
+  aggregate lifespan.
 
   By default an aggregate instance process will run indefinitely once started.
-  You can control this by implementing the `Commanded.Aggregates.AggregateLifespan` behaviour in a module.
+  You can control this by implementing the
+  `Commanded.Aggregates.AggregateLifespan` behaviour in a module and configuring
+  it in your router.
 
-  The `c:after_command/1` function is called after each command executes.
-  The returned timeout value is used to shutdown the aggregate process if no other messages are recevied.
+  After a command successfully executes, and creates at least one domain event,
+  the `c:after_event/1` function is called passing the last created event. The
+  returned inactivity timeout value is used to shutdown the aggregate process if
+  no other messages are received.
 
-  Return `:infinity` to prevent the aggregate instance from shutting down.
-  """
+  ## Supported return values
 
-  @type command :: struct()
-  @type timeout_or_action :: timeout() | :infinity | :hibernate
+    - Non-negative integer - specify an inactivity timeout, in millisconds.
+    - `:infinity` - prevent the aggregate instance from shutting down.
+    - `:hibernate` - send the process into hibernation.
+    - `:stop` - immediately shutdown the aggregate process.
 
-  @doc """
-  After specified timeout aggregate process will be stoped.
+  ### Hibernation
+
+  A hibernated process will continue its loop once a message is in its message
+  queue. Hibernating an aggregate causes garbage collection and minimises the
+  memory used by the process. Hibernating should not be used aggressively as too
+  much time could be spent garbage collecting.
+
+  ## Example
+
+  Define a module that implements the `Commanded.Aggregates.AggregateLifespan`
+  behaviour:
+
+      defmodule BankAccountLifespan do
+        @behaviour Commanded.Aggregates.AggregateLifespan
+
+        def after_event(%BankAccountOpened{}), do: :infinity
+        def after_event(%MoneyDeposited{}), do: 60_000
+        def after_event(%BankAccountClosed{}), do: :stop
+      end
 
-  Timeout begins at start of processing a command.
+  Then specify the module as the `lifespan` option when registering
+  the applicable commands in your router:
+
+      defmodule BankRouter do
+        use Commanded.Commands.Router
+
+        dispatch [OpenAccount, CloseAccount],
+          to: BankAccount,
+          lifespan: BankAccountLifespan,
+          identity: :account_number
+      end
+
+  """
+
+  @doc """
+  Aggregate process will be stopped after specified inactivity timeout unless
+  `:infinity`, `:hibernate`, or `:stop` are returned.
   """
-  @callback after_command(command) :: timeout_or_action
+  @callback after_event(event :: struct()) :: timeout() | :hibernate | :stop
 end
diff --git a/lib/commanded/aggregates/default_lifespan.ex b/lib/commanded/aggregates/default_lifespan.ex
index 1f3af52..8d0bfde 100644
--- a/lib/commanded/aggregates/default_lifespan.ex
+++ b/lib/commanded/aggregates/default_lifespan.ex
@@ -1,14 +1,16 @@
 defmodule Commanded.Aggregates.DefaultLifespan do
   @moduledoc """
-  The default implementation of the `Commanded.Aggregates.AggregateLifespan` behaviour.
+  The default implementation of the `Commanded.Aggregates.AggregateLifespan`
+  behaviour.
 
-  It will ensure that an aggregate instance process runs indefinitely once started.
+  It will ensure that an aggregate instance process runs indefinitely once
+  started.
   """
 
   @behaviour Commanded.Aggregates.AggregateLifespan
 
   @doc """
-  Aggregate will run indefinitely once started
+  Aggregate will run indefinitely once started.
   """
-  def after_command(_command), do: :infinity
+  def after_event(_event), do: :infinity
 end
diff --git a/lib/commanded/commands/router.ex b/lib/commanded/commands/router.ex
index e7462d0..c40dbe4 100644
--- a/lib/commanded/commands/router.ex
+++ b/lib/commanded/commands/router.ex
@@ -263,17 +263,35 @@ defmodule Commanded.Commands.Router do
   do
     quote location: :keep do
       if Enum.member?(@registered_commands, unquote(command_module)) do
-        raise "duplicate command registration for: #{inspect unquote(command_module)}"
+        raise ArgumentError,
+          message: "Command `#{inspect unquote(command_module)}` has already been registered in router `#{inspect __MODULE__}`"
       end
 
       # sanity check the configured modules exist
+      ensure_module_exists(unquote(aggregate))
       ensure_module_exists(unquote(command_module))
       ensure_module_exists(unquote(handler))
-      ensure_module_exists(unquote(aggregate))
 
-      handler_functions = unquote(handler).__info__(:functions)
-      unless Keyword.get(handler_functions, unquote(function)) == 2 do
-        raise "command handler #{inspect unquote(handler)} does not define a function: #{unquote(function)}/2"
+      case unquote(lifespan) do
+        nil ->
+          :ok
+
+        module when is_atom(module) ->
+          ensure_module_exists(unquote(lifespan))
+
+          unless function_exported?(unquote(lifespan), :after_event, 1) do
+            raise ArgumentError,
+              message: "Aggregate lifespan `#{inspect unquote(lifespan)}` does not define a callback function: `after_event/1`"
+          end
+
+        invalid ->
+          raise ArgumentError,
+            message: "Invalid `lifespan` configured for #{inspect unquote(aggregate)}: #{inspect invalid}"
+      end
+
+      unless function_exported?(unquote(handler), unquote(function), 2) do
+        raise ArgumentError,
+          message: "Command handler `#{inspect unquote(handler)}` does not define a `#{unquote(function)}/2` function"
       end
 
       @registered_commands [unquote(command_module) | @registered_commands]


Test patch:
diff --git a/test/aggregates/aggregate_lifespan_test.exs b/test/aggregates/aggregate_lifespan_test.exs
index 2a0d269..3df324a 100644
--- a/test/aggregates/aggregate_lifespan_test.exs
+++ b/test/aggregates/aggregate_lifespan_test.exs
@@ -3,18 +3,22 @@ defmodule Commanded.Aggregates.AggregateLifespanTest do
 
   alias Commanded.Aggregates.BankRouter
   alias Commanded.ExampleDomain.BankAccount
+
   alias Commanded.ExampleDomain.BankAccount.Commands.{
+    CloseAccount,
     OpenAccount,
     DepositMoney,
-    WithdrawMoney,
+    WithdrawMoney
   }
+
   alias Commanded.Registration
 
   describe "aggregate started" do
     setup do
-      aggregate_uuid = UUID.uuid4
+      aggregate_uuid = UUID.uuid4()
 
-      {:ok, ^aggregate_uuid} = Commanded.Aggregates.Supervisor.open_aggregate(BankAccount, aggregate_uuid)
+      {:ok, ^aggregate_uuid} =
+        Commanded.Aggregates.Supervisor.open_aggregate(BankAccount, aggregate_uuid)
 
       pid = Registration.whereis_name({BankAccount, aggregate_uuid})
       ref = Process.monitor(pid)
@@ -25,22 +29,64 @@ defmodule Commanded.Aggregates.AggregateLifespanTest do
     test "should shutdown after timeout", %{aggregate_uuid: aggregate_uuid, ref: ref} do
       :ok = BankRouter.dispatch(%OpenAccount{account_number: aggregate_uuid, initial_balance: 10})
 
-      assert_receive {:DOWN, ^ref, :process, _, :normal}, 10
+      assert_receive {:DOWN, ^ref, :process, _, :normal}
     end
 
-    test "should not shutdown if next command executed", %{aggregate_uuid: aggregate_uuid, ref: ref} do
+    test "should not shutdown if another command executed", %{
+      aggregate_uuid: aggregate_uuid,
+      ref: ref
+    } do
       :ok = BankRouter.dispatch(%OpenAccount{account_number: aggregate_uuid, initial_balance: 10})
       :ok = BankRouter.dispatch(%DepositMoney{account_number: aggregate_uuid, amount: 10})
 
       refute_receive {:DOWN, ^ref, :process, _, :normal}, 10
-      assert_receive {:DOWN, ^ref, :process, _, :normal}, 30
+      assert_receive {:DOWN, ^ref, :process, _, :normal}
     end
 
-    test "should use default lifespan when it's not specified'", %{aggregate_uuid: aggregate_uuid, ref: ref} do
+    test "should use default lifespan when it's not specified'", %{
+      aggregate_uuid: aggregate_uuid,
+      ref: ref
+    } do
       :ok = BankRouter.dispatch(%OpenAccount{account_number: aggregate_uuid, initial_balance: 10})
       :ok = BankRouter.dispatch(%WithdrawMoney{account_number: aggregate_uuid, amount: 10})
 
       refute_receive {:DOWN, ^ref, :process, _, :normal}, 30
     end
+
+    test "should stop process when requested" , %{
+      aggregate_uuid: aggregate_uuid,
+      ref: ref
+    } do
+      :ok = BankRouter.dispatch(%OpenAccount{account_number: aggregate_uuid, initial_balance: 10})
+      :ok = BankRouter.dispatch(%CloseAccount{account_number: aggregate_uuid})
+
+      assert_receive {:DOWN, ^ref, :process, _, :normal}
+    end
+  end
+
+  describe "deprecated `after_command/1` callback" do
+    test "should fail to compile when missing `after_event/1` function" do
+      assert_raise ArgumentError, "Aggregate lifespan `BankAccountLifespan` does not define a callback function: `after_event/1`", fn ->
+        Code.eval_string """
+          alias Commanded.ExampleDomain.BankAccount
+          alias Commanded.ExampleDomain.BankAccount.Commands.{OpenAccount, DepositMoney}
+
+          defmodule BankAccountLifespan do
+            def after_command(%OpenAccount{}), do: 5
+            def after_command(%DepositMoney{}), do: 20
+            def after_command(_), do: :infinity
+          end
+
+          defmodule BankRouter do
+            use Commanded.Commands.Router
+
+            dispatch [OpenAccount],
+              to: BankAccount,
+              lifespan: BankAccountLifespan,
+              identity: :account_number
+          end
+        """
+      end
+    end
   end
 end
diff --git a/test/aggregates/support/bank_account_lifespan.ex b/test/aggregates/support/bank_account_lifespan.ex
index 714e749..7b366b2 100644
--- a/test/aggregates/support/bank_account_lifespan.ex
+++ b/test/aggregates/support/bank_account_lifespan.ex
@@ -1,13 +1,16 @@
 defmodule Commanded.Aggregates.BankAccountLifespan do
   @moduledoc false
+
   @behaviour Commanded.Aggregates.AggregateLifespan
 
-  alias Commanded.ExampleDomain.BankAccount.Commands.{
-    OpenAccount,
-    DepositMoney,
+  alias Commanded.ExampleDomain.BankAccount.Events.{
+    BankAccountClosed,
+    BankAccountOpened,
+    MoneyDeposited
   }
 
-  def after_command(%OpenAccount{}), do: 5
-  def after_command(%DepositMoney{}), do: 20
-  def after_command(_), do: :infinity
+  def after_event(%BankAccountOpened{}), do: 5
+  def after_event(%MoneyDeposited{}), do: 20
+  def after_event(%BankAccountClosed{}), do: :stop
+  def after_event(_event), do: :infinity
 end
diff --git a/test/aggregates/support/bank_router.ex b/test/aggregates/support/bank_router.ex
index 89c9ee1..b1cf2ba 100644
--- a/test/aggregates/support/bank_router.ex
+++ b/test/aggregates/support/bank_router.ex
@@ -3,16 +3,19 @@ defmodule Commanded.Aggregates.BankRouter do
   use Commanded.Commands.Router
 
   alias Commanded.Aggregates.BankAccountLifespan
+
   alias Commanded.ExampleDomain.{
     BankAccount,
     OpenAccountHandler,
     DepositMoneyHandler,
-    WithdrawMoneyHandler,
+    WithdrawMoneyHandler
   }
+
   alias BankAccount.Commands.{
-    OpenAccount,
+    CloseAccount,
     DepositMoney,
-    WithdrawMoney,
+    OpenAccount,
+    WithdrawMoney
   }
 
   dispatch OpenAccount,
@@ -31,4 +34,10 @@ defmodule Commanded.Aggregates.BankRouter do
     to: WithdrawMoneyHandler,
     aggregate: BankAccount,
     identity: :account_number
+
+  dispatch CloseAccount,
+    to: OpenAccountHandler,
+    aggregate: BankAccount,
+    lifespan: BankAccountLifespan,
+    identity: :account_number
 end
diff --git a/test/commands/routing_commands_test.exs b/test/commands/routing_commands_test.exs
index ad43e51..60079bb 100644
--- a/test/commands/routing_commands_test.exs
+++ b/test/commands/routing_commands_test.exs
@@ -123,7 +123,7 @@ defmodule Commanded.Commands.RoutingCommandsTest do
   end
 
   test "should prevent duplicate registrations for a command" do
-    assert_raise RuntimeError, "duplicate command registration for: Commanded.ExampleDomain.BankAccount.Commands.OpenAccount", fn ->
+    assert_raise ArgumentError, "Command `Commanded.ExampleDomain.BankAccount.Commands.OpenAccount` has already been registered in router `DuplicateRouter`", fn ->
       Code.eval_string """
         alias Commanded.ExampleDomain.BankAccount
         alias Commanded.ExampleDomain.BankAccount.Commands.OpenAccount
@@ -143,7 +143,7 @@ defmodule Commanded.Commands.RoutingCommandsTest do
   end
 
   test "should prevent registration for a command handler without a `handle/2` function" do
-    assert_raise RuntimeError, "command handler InvalidHandler does not define a function: handle/2", fn ->
+    assert_raise ArgumentError, "Command handler `InvalidHandler` does not define a `handle/2` function", fn ->
       Code.eval_string """
         alias Commanded.ExampleDomain.BankAccount
         alias Commanded.ExampleDomain.BankAccount.Commands.OpenAccount
diff --git a/test/example_domain/bank_account/open_account_handler.ex b/test/example_domain/bank_account/open_account_handler.ex
index 5e4430d..b43f4a3 100644
--- a/test/example_domain/bank_account/open_account_handler.ex
+++ b/test/example_domain/bank_account/open_account_handler.ex
@@ -1,17 +1,16 @@
 defmodule Commanded.ExampleDomain.OpenAccountHandler do
   @moduledoc false
+
   alias Commanded.ExampleDomain.BankAccount
-  alias Commanded.ExampleDomain.BankAccount.Commands.{OpenAccount,CloseAccount}
+  alias Commanded.ExampleDomain.BankAccount.Commands.{OpenAccount, CloseAccount}
 
   @behaviour Commanded.Commands.Handler
 
   def handle(%BankAccount{} = aggregate, %OpenAccount{} = open_account) do
-    aggregate
-    |> BankAccount.open_account(open_account)
+    BankAccount.open_account(aggregate, open_account)
   end
 
   def handle(%BankAccount{} = aggregate, %CloseAccount{} = close_account) do
-    aggregate
-    |> BankAccount.close_account(close_account)
+    BankAccount.close_account(aggregate, close_account)
   end
 end