Close SSH session before exiting

lib/producer/core/cli.rb

@@ -54,6 +54,7 @@ module Producer
 
       def run
         worker.process load_recipe.tasks
+        env.cleanup
       end
 
       def load_recipe

lib/producer/core/env.rb

@@ -44,6 +44,10 @@ module Producer
       def dry_run?
         @dry_run
       end
+
+      def cleanup
+        remote.cleanup
+      end
     end
   end
 end

lib/producer/core/remote.rb

@@ -2,6 +2,7 @@ module Producer
   module Core
     class Remote
       attr_reader :hostname
+      attr_writer :session
 
       def initialize(hostname)
         @hostname = hostname
@@ -43,6 +44,10 @@ module Producer
       def environment
         Environment.new_from_string(execute 'env')
       end
+
+      def cleanup
+        session.close if @session
+      end
     end
   end
 end

spec/producer/core/cli_spec.rb

@@ -172,6 +172,11 @@ module Producer::Core
           .with([an_instance_of(Task), an_instance_of(Task)])
         cli.run
       end
+
+      it 'cleans up the env' do
+        expect(cli.env).to receive :cleanup
+        cli.run
+      end
     end
 
     describe '#load_recipe' do

spec/producer/core/env_spec.rb

@@ -157,5 +157,12 @@ module Producer::Core
         expect(env).to be_dry_run
       end
     end
+
+    describe '#cleanup' do
+      it 'cleans up the remote' do
+        expect(env.remote).to receive :cleanup
+        env.cleanup
+      end
+    end
   end
 end

spec/producer/core/remote_spec.rb

@@ -136,5 +136,14 @@ module Producer::Core
         expect(remote.environment['FOO']).to eq 'bar'
       end
     end
+
+    describe '#cleanup' do
+      before { remote.session = double 'session' }
+
+      it 'closes the session' do
+        expect(remote.session).to receive :close
+        remote.cleanup
+      end
+    end
   end
 end