Handle child spec options in Chaotic workers/supervisors

Author: fishcakez

lib/chaos_spawn/chaotic/chaotic_supervisor.ex

@@ -4,17 +4,9 @@ defmodule ChaosSpawn.Chaotic.ChaoticSupervisor do
   register with ChaosSpawn as a killable process.
 
   """
-  alias Supervisor.Spec, as: OriginalSupervisor
   alias ChaosSpawn.Chaotic.Supervisor.Wrapper
 
   def supervisor(module, args, opts \\ []) do
-    start_link_function = opts |> Keyword.get(:function, :start_link)
-    id = opts |> Keyword.get(:id, module)
-    OriginalSupervisor.supervisor(
-      Wrapper,
-      [module, start_link_function, args],
-      id: id,
-      function: :start_link_wrapper
-    )
+    Wrapper.child_spec(:supervisor, module, args, opts)
   end
 end

lib/chaos_spawn/chaotic/chaotic_worker.ex

@@ -3,18 +3,9 @@ defmodule ChaosSpawn.Chaotic.ChaoticWorker do
   Provides worker/2 and worker/3 that wrap around Supervisor.Spec's worker
   functions and registered any spawned pids with ChaosSpawn's process killer.
   """
-  alias Supervisor.Spec, as: OriginalSupervisor
   alias ChaosSpawn.Chaotic.Supervisor.Wrapper
 
   def worker(module, args, opts \\ []) do
-    start_link_function = opts |> Keyword.get(:function, :start_link)
-    id = opts |> Keyword.get(:id, module)
-    OriginalSupervisor.worker(
-      Wrapper,
-      [module, start_link_function, args],
-      id: id,
-      function: :start_link_wrapper
-    )
+    Wrapper.child_spec(:worker, module, args, opts)
   end
-
 end

lib/chaos_spawn/chaotic/supervisor/wrapper.ex

@@ -19,6 +19,17 @@ defmodule ChaosSpawn.Chaotic.Supervisor.Wrapper do
       |> register_unless_skipped(module, skipped)
   end
 
+  @doc false
+  def child_spec(type, module, args, opts) do
+    start_link_function = opts |> Keyword.get(:function, :start_link)
+    args = [module, start_link_function, args]
+    opts = opts
+      |> Keyword.put(:function, :start_link_wrapper)
+      |> Keyword.put_new(:id, module)
+      |> Keyword.put_new(:modules, [module])
+    apply(Supervisor.Spec, type, [__MODULE__, args, opts])
+  end
+
   defp register_unless_skipped(result, module, skipped_modules) do
     if (skip_module?(skipped_modules, module)) do
       result

test/chaos_spawn/chaotic/supervisor_test.exs

@@ -13,7 +13,7 @@ defmodule Chaotic.SupervisorTest do
       :permanent,
       :infinity,
       :supervisor,
-      [Wrapper]
+      [ModuleToCall]
     }
 
     assert ChaoticSupervisor.supervisor(ModuleToCall, args) == expected

test/chaos_spawn/chaotic/worker_test.exs

@@ -13,7 +13,7 @@ defmodule Chaotic.WorkerTest do
       :permanent,
       5000,
       :worker,
-      [Wrapper]
+      [ModuleToCall]
     }
 
     assert ChaoticWorker.worker(ModuleToCall, args) == expected
@@ -29,10 +29,31 @@ defmodule Chaotic.WorkerTest do
       :permanent,
       5000,
       :worker,
-      [Wrapper]
+      [ModuleToCall]
     }
 
     worker = ChaoticWorker.worker(ModuleToCall, args, function: my_start_func)
     assert worker == expected
   end
+
+  test "worker/3 sets :id, :restart, :shutdown, :modules" do
+    args = [:arg_one, :arg_two]
+
+    expected = {
+      :special_id,
+      {Wrapper, :start_link_wrapper, [ModuleToCall, :start_link, args]},
+      :transient,
+      :brutal_kill,
+      :worker,
+      :dynamic
+    }
+
+    worker = ChaoticWorker.worker(ModuleToCall, args,
+      id: :special_id,
+      restart: :transient,
+      shutdown: :brutal_kill,
+      modules: :dynamic
+    )
+    assert worker == expected
+  end
 end