Updates to ensure connections are closed in ConnectionManager and tests

README.md

@@ -137,15 +137,14 @@ Creating faktory_worker_test_tls      ... done
 Creating faktory_worker_password_test ... done
 ```
 
+Faktory has a free, open-source solution and an enterprise edition. By default, tests for the enterprise edition of Faktory are excluded:
 
-Faktory have free open-source solution and enterprise edition.
-
-If you don't have enterprise license then tests will fail on enterprise features (batching operations etc). In this case you can exclude them by tag `:enterprise`
 ```sh
-$ mix test --exclude enterprise
+$ mix test
 ```
 
-If you are enterprise user all tests should pass
+If you wish to run the enterprise tests, you may include them, like so:
+
 ```sh
-$ mix test
+$ mix test --include enterprise
 ```

lib/faktory_worker/connection_manager.ex

@@ -32,14 +32,14 @@ defmodule FaktoryWorker.ConnectionManager do
           {Connection.response(), ConnectionManager.t()}
   def send_command(state, command, allow_retry \\ true) do
     case try_send_command(state, command) do
-      {{:error, reason}, _} when reason in @connection_errors ->
+      {{:error, reason}, state} when reason in @connection_errors ->
         # Close dangling port
         close_connection(state)
         error = {:error, "Failed to connect to Faktory"}
         state = %{state | conn: nil}
 
         if allow_retry,
-          do: send_command(%{state | conn: nil}, command, false),
+          do: send_command(state, command, false),
           else: {error, state}
 
       # Handle errors from Faktory that should not be tried again

test/faktory_worker/connection_manager_test.exs

@@ -97,13 +97,10 @@ defmodule FaktoryWorker.ConnectionManagerTest do
     test "should unset the connection when there is a socket failure" do
       connection_mox()
 
-      expect(FaktoryWorker.SocketMock, :send, fn _, "INFO\r\n" ->
-        {:error, :closed}
-      end)
-
-      expect(FaktoryWorker.SocketMock, :connect, fn _, _, _ ->
-        {:error, :econnrefused}
-      end)
+      FaktoryWorker.SocketMock
+      |> expect(:send, fn _, "INFO\r\n" -> {:error, :closed} end)
+      |> expect(:connect, fn _, _, _ -> {:error, :econnrefused} end)
+      |> expect(:close, fn _ -> :ok end)
 
       opts = [socket_handler: FaktoryWorker.SocketMock]
       state = ConnectionManager.new(opts)
@@ -161,13 +158,10 @@ defmodule FaktoryWorker.ConnectionManagerTest do
 
       connection_mox()
 
-      expect(FaktoryWorker.SocketMock, :send, fn _, _ ->
-        :ok
-      end)
-
-      expect(FaktoryWorker.SocketMock, :recv, fn _ ->
-        {:ok, "+OK\r\n"}
-      end)
+      FaktoryWorker.SocketMock
+      |> expect(:send, fn _, _ -> :ok end)
+      |> expect(:recv, fn _ -> {:ok, "+OK\r\n"} end)
+      |> expect(:close, fn _ -> :ok end)
 
       opts = [socket_handler: FaktoryWorker.SocketMock]
 

test/faktory_worker/worker/heartbeat_server_test.exs

@@ -241,17 +241,11 @@ defmodule FaktoryWorker.Worker.HeartbeatServerTest do
 
       worker_connection_mox()
 
-      expect(FaktoryWorker.SocketMock, :send, fn _, ^beat_command ->
-        :ok
-      end)
-
-      expect(FaktoryWorker.SocketMock, :recv, fn _ ->
-        {:error, :closed}
-      end)
-
-      expect(FaktoryWorker.SocketMock, :connect, fn _, _, _ ->
-        {:error, :econnrefused}
-      end)
+      FaktoryWorker.SocketMock
+      |> expect(:send, fn _, ^beat_command -> :ok end)
+      |> expect(:recv, fn _ -> {:error, :closed} end)
+      |> expect(:connect, fn _, _, _ -> {:error, :econnrefused} end)
+      |> expect(:close, fn _ -> :ok end)
 
       opts = [socket_handler: FaktoryWorker.SocketMock, is_worker: true, process_wid: process_wid]
 

test/faktory_worker/worker_test.exs

@@ -792,24 +792,18 @@ defmodule FaktoryWorker.WorkerTest do
 
       worker_connection_mox()
 
-      expect(FaktoryWorker.SocketMock, :send, fn _, ^ack_command ->
-        :ok
-      end)
-
-      expect(FaktoryWorker.SocketMock, :recv, fn _ ->
-        {:error, :closed}
-      end)
+      FaktoryWorker.SocketMock
+      |> expect(:send, fn _, ^ack_command -> :ok end)
+      |> expect(:recv, fn _ -> {:error, :closed} end)
+      |> expect(:close, fn _ -> :ok end)
 
       # the connection manager retries a failed request once
       worker_connection_mox()
 
-      expect(FaktoryWorker.SocketMock, :send, fn _, ^ack_command ->
-        :ok
-      end)
-
-      expect(FaktoryWorker.SocketMock, :recv, fn _ ->
-        {:error, :closed}
-      end)
+      FaktoryWorker.SocketMock
+      |> expect(:send, fn _, ^ack_command -> :ok end)
+      |> expect(:recv, fn _ -> {:error, :closed} end)
+      |> expect(:close, fn _ -> :ok end)
 
       opts = [
         process_wid: Random.process_wid(),
@@ -858,24 +852,18 @@ defmodule FaktoryWorker.WorkerTest do
 
       worker_connection_mox()
 
-      expect(FaktoryWorker.SocketMock, :send, fn _, ^fail_command ->
-        :ok
-      end)
-
-      expect(FaktoryWorker.SocketMock, :recv, fn _ ->
-        {:error, :closed}
-      end)
+      FaktoryWorker.SocketMock
+      |> expect(:send, fn _, ^fail_command -> :ok end)
+      |> expect(:recv, fn _ -> {:error, :closed} end)
+      |> expect(:close, fn _ -> :ok end)
 
       # the connection manager retries a failed request one more time
       worker_connection_mox()
 
-      expect(FaktoryWorker.SocketMock, :send, fn _, ^fail_command ->
-        :ok
-      end)
-
-      expect(FaktoryWorker.SocketMock, :recv, fn _ ->
-        {:error, :closed}
-      end)
+      FaktoryWorker.SocketMock
+      |> expect(:send, fn _, ^fail_command -> :ok end)
+      |> expect(:recv, fn _ -> {:error, :closed} end)
+      |> expect(:close, fn _ -> :ok end)
 
       opts = [
         process_wid: Random.process_wid(),