From d400b5db68aad08344396114e6d29360ba6dadbc Mon Sep 17 00:00:00 2001
From: seenmyfate <seenmyfate@gmail.com>
Date: Fri, 31 May 2013 16:36:20 +0100
Subject: [PATCH] Support filtering of roles

    # config/deploy/stage.rb
    server 'example1.com', roles: %w{web app db}, active: :true
    server 'example2.com', roles: %w{web app db}

    on :app, filter: { active: true } do
      # do things on just example1.com
    end
---
 lib/capistrano/configuration.rb               | 17 +---
 lib/capistrano/configuration/server.rb        | 28 ++++--
 lib/capistrano/configuration/servers.rb       | 68 ++++++++++++--
 lib/capistrano/dsl/env.rb                     |  5 +-
 lib/capistrano/i18n.rb                        |  1 +
 .../capistrano/configuration/server_spec.rb   |  4 +-
 .../capistrano/configuration/servers_spec.rb  | 89 ++++++++++++++++---
 spec/lib/capistrano/configuration_spec.rb     |  4 +-
 spec/lib/capistrano/dsl/env_spec.rb           | 73 ---------------
 9 files changed, 171 insertions(+), 118 deletions(-)

diff --git a/lib/capistrano/configuration.rb b/lib/capistrano/configuration.rb
index 37adb06c..7d5ed470 100644
--- a/lib/capistrano/configuration.rb
+++ b/lib/capistrano/configuration.rb
@@ -33,23 +33,12 @@ module Capistrano
       servers.add_role(name, hosts)
     end
 
-    def server(name, properties = {})
+    def server(name, properties={})
       servers.add_host(name, properties)
     end
 
-    def roles_for(names, options = {})
-      servers.fetch_roles(names).tap do |list|
-        if filter = options.delete(:filter) || options.delete(:select)
-          if filter.respond_to?(:call)
-            list.select!(&filter)
-          else
-            list.select! { |s| s.properties.send(filter) }
-          end
-          if list.empty?
-            raise "Your filter #{filter} would remove all matching servers!"
-          end
-        end
-      end
+    def roles_for(names)
+      servers.roles_for(names)
     end
 
     def primary(role)
diff --git a/lib/capistrano/configuration/server.rb b/lib/capistrano/configuration/server.rb
index b29f958b..9b647ab8 100644
--- a/lib/capistrano/configuration/server.rb
+++ b/lib/capistrano/configuration/server.rb
@@ -8,7 +8,7 @@ module Capistrano
       end
 
       def add_role(role)
-        roles << role.to_sym
+        roles.add role.to_sym
       end
 
       def has_role?(role)
@@ -20,25 +20,37 @@ module Capistrano
       end
 
       def roles
-        properties.roles ||= Set.new
+        properties.cap_roles ||= Set.new
       end
 
-      def primary
-        self if properties.primary
+      def primary?
+        self if fetch(:primary?)
       end
 
       def with(properties)
-        properties.each { |property, value| add_property(property, value) }
+        properties.each { |key, value| add_property(key, value) }
         self
       end
 
+      def fetch(key)
+        properties.send(key)
+      end
+
+      def set(key, value)
+        properties.send(:"#{key}=", value) unless set?(key)
+      end
+
+      def set?(key)
+        properties.respond_to?(key)
+      end
+
       private
 
-      def add_property(property, value)
-        if property.to_sym == :role
+      def add_property(key, value)
+        if key.to_sym == :roles
           add_roles(value)
         else
-          properties.send(:"#{property}=", value) unless properties.respond_to?(property)
+          set(key, value)
         end
       end
 
diff --git a/lib/capistrano/configuration/servers.rb b/lib/capistrano/configuration/servers.rb
index 8ee81392..7aa4cf9a 100644
--- a/lib/capistrano/configuration/servers.rb
+++ b/lib/capistrano/configuration/servers.rb
@@ -9,16 +9,17 @@ module Capistrano
       end
 
       def add_role(role, hosts)
-        Array(hosts).each { |host| add_host(host, role: role) }
+        Array(hosts).each { |host| add_host(host, roles: role) }
       end
 
-      def fetch_roles(names)
-        roles_for(names)
+      def roles_for(*names)
+        options = extract_options(names)
+        fetch_roles(*names, options)
       end
 
       def fetch_primary(role)
         hosts = fetch(role)
-        hosts.find(&:primary) || hosts.first
+        hosts.find(&:primary?) || hosts.first
       end
 
       def each
@@ -28,24 +29,75 @@ module Capistrano
       private
 
       def server(host)
-        servers.find { |server| server.matches?(host) } || Server.new(host)
+        if host.is_a? Server
+          host
+        else
+          servers.find { |server| server.matches?(host) } || Server.new(host)
+        end
       end
 
       def fetch(role)
         servers.find_all { |server| server.has_role? role}
       end
 
-      def roles_for(names)
+      def fetch_roles(names, options)
         if Array(names).map(&:to_sym).include?(:all)
-          servers
+          filter(servers, options)
         else
-          Array(names).flat_map { |name| fetch name }.uniq
+          role_servers = Array(names).flat_map { |name| fetch name }.uniq
+          filter(role_servers, options)
         end
       end
 
+      def filter(servers, options)
+        Filter.new(servers, options).filtered_servers
+      end
+
       def servers
         @servers ||= Set.new
       end
+
+      def extract_options(array)
+        array.last.is_a?(::Hash) ? array.pop : {}
+      end
+
+      class Filter
+        def initialize(servers, options)
+          @servers, @options = servers, options
+        end
+
+        def filtered_servers
+          if servers_with_filter.any?
+            servers_with_filter
+          else
+            fail I18n.t(:filter_removes_all_servers)
+          end
+        end
+
+        private
+        attr_reader :options, :servers
+
+        def servers_with_filter
+          @servers_with_filter ||= servers.select(&filter)
+        end
+
+        def filter_option
+          options[:filter] || options[:select] || all
+        end
+
+        def filter
+          if filter_option.respond_to?(:call)
+            filter_option
+          else
+            lambda { |server| server.fetch(filter_option) }
+          end
+        end
+
+        def all
+          lambda { |server| :all }
+        end
+
+      end
     end
   end
 end
diff --git a/lib/capistrano/dsl/env.rb b/lib/capistrano/dsl/env.rb
index 76f74863..68d095a5 100644
--- a/lib/capistrano/dsl/env.rb
+++ b/lib/capistrano/dsl/env.rb
@@ -31,6 +31,10 @@ module Capistrano
         env.role(name, servers)
       end
 
+      def server(name, properties={})
+        env.server(name, properties)
+      end
+
       def roles(*names)
         env.roles_for(names)
       end
@@ -54,4 +58,3 @@ module Capistrano
     end
   end
 end
-
diff --git a/lib/capistrano/i18n.rb b/lib/capistrano/i18n.rb
index ab3c94aa..a29542b8 100644
--- a/lib/capistrano/i18n.rb
+++ b/lib/capistrano/i18n.rb
@@ -18,6 +18,7 @@ en = {
   mirror_exists:               "The repository mirror is at %{at}",
   revision_log_message:        'Branch %{branch} deployed as release %{release} by %{user}',
   rollback_log_message:        '%{user} rolled back to release %{release}',
+  filter_removes_all_servers:  'Your filter %{filter} would remove all matching servers',
   console: {
     welcome:  'capistrano console - enter command to execute on %{stage}',
     bye:      'bye'
diff --git a/spec/lib/capistrano/configuration/server_spec.rb b/spec/lib/capistrano/configuration/server_spec.rb
index e77b8e3b..6fdd0433 100644
--- a/spec/lib/capistrano/configuration/server_spec.rb
+++ b/spec/lib/capistrano/configuration/server_spec.rb
@@ -52,10 +52,10 @@ module Capistrano
       end
 
       describe 'identifying as primary' do
-        subject { server.primary }
+        subject { server.primary? }
         context 'server is primary' do
           before do
-            server.properties.primary = true
+            server.set(:primary?, true)
           end
           it 'returns self' do
             expect(subject).to eq server
diff --git a/spec/lib/capistrano/configuration/servers_spec.rb b/spec/lib/capistrano/configuration/servers_spec.rb
index 781a92b1..a9ffe506 100644
--- a/spec/lib/capistrano/configuration/servers_spec.rb
+++ b/spec/lib/capistrano/configuration/servers_spec.rb
@@ -6,11 +6,28 @@ module Capistrano
       let(:servers) { Servers.new }
 
       describe 'adding a role' do
-        subject { servers.add_role(:app, %w{1 2}) }
 
         it 'adds two new server instances' do
-          expect{subject}.to change{servers.count}.from(0).to(2)
+          expect{servers.add_role(:app, %w{1 2})}.
+            to change{servers.count}.from(0).to(2)
         end
+
+        it 'handles de-duplification within roles' do
+          servers.add_role(:app, %w{1})
+          servers.add_role(:app, %w{1})
+          expect(servers.count).to eq 1
+        end
+
+        it 'accepts instances of server objects' do
+          servers.add_role(:app, [Capistrano::Configuration::Server.new('example.net'), 'example.com'])
+          expect(servers.roles_for(:app).length).to eq 2
+        end
+
+        it 'accepts non-enumerable types' do
+          servers.add_role(:app, '1')
+          expect(servers.roles_for(:app).count).to eq 1
+        end
+
       end
 
       describe 'adding a role to an existing server' do
@@ -22,6 +39,7 @@ module Capistrano
         it 'adds new roles to existing servers' do
           expect(servers.count).to eq 2
         end
+
       end
 
       describe 'collecting server roles' do
@@ -35,21 +53,21 @@ module Capistrano
         end
 
         it 'returns an array of the roles' do
-          expect(servers.fetch_roles([:app]).collect(&:roles)).to eq [app, web_app, web_app]
-          expect(servers.fetch_roles([:web]).collect(&:roles)).to eq [web_app, web_app, web]
+          expect(servers.roles_for([:app]).collect(&:roles)).to eq [app, web_app, web_app]
+          expect(servers.roles_for([:web]).collect(&:roles)).to eq [web_app, web_app, web]
         end
       end
 
       describe 'finding the primary server' do
         it 'takes the first server if none have the primary property' do
           servers.add_role(:app, %w{1 2})
-          servers.fetch_primary(:app).hostname.should == "1"
+          servers.fetch_primary(:app).hostname.should == '1'
         end
 
         it 'takes the first server with the primary have the primary flag' do
           servers.add_role(:app, %w{1 2})
-          servers.add_host('2', primary: true)
-          servers.fetch_primary(:app).hostname.should == "2"
+          servers.add_host('2', primary?: true)
+          servers.fetch_primary(:app).hostname.should == '2'
         end
       end
 
@@ -60,21 +78,70 @@ module Capistrano
         end
 
         it 'returns the correct app servers' do
-          expect(servers.fetch_roles([:app]).map(&:hostname)).to eq %w{1 2}
+          expect(servers.roles_for([:app]).map(&:hostname)).to eq %w{1 2}
         end
 
         it 'returns the correct web servers' do
-          expect(servers.fetch_roles([:web]).map(&:hostname)).to eq %w{2 3}
+          expect(servers.roles_for([:web]).map(&:hostname)).to eq %w{2 3}
         end
 
         it 'returns the correct app and web servers' do
-          expect(servers.fetch_roles([:app, :web]).map(&:hostname)).to eq %w{1 2 3}
+          expect(servers.roles_for([:app, :web]).map(&:hostname)).to eq %w{1 2 3}
         end
 
         it 'returns all servers' do
-          expect(servers.fetch_roles([:all]).map(&:hostname)).to eq %w{1 2 3}
+          expect(servers.roles_for([:all]).map(&:hostname)).to eq %w{1 2 3}
         end
       end
+
+      describe 'adding a server' do
+
+        before do
+          servers.add_host('1', roles: [:app, 'web'], test: :value)
+        end
+
+        it 'can create a server with properties' do
+          expect(servers.roles_for([:app]).first.hostname).to eq '1'
+          expect(servers.roles_for([:web]).first.hostname).to eq '1'
+          expect(servers.roles_for([:all]).first.properties.test).to eq :value
+        end
+
+      end
+
+      describe '#roles' do
+
+        before do
+          servers.add_host('1', roles: :app, active: true)
+          servers.add_host('2', roles: :app)
+        end
+
+        it 'raises if the filter would remove all matching hosts' do
+          I18n.expects(:t).with(:filter_removes_all_servers)
+          expect { servers.roles_for(:app, select: :inactive) }.to raise_error
+        end
+
+        it 'can filter hosts by properties on the host object using symbol as shorthand' do
+          expect(servers.roles_for(:app, filter: :active).length).to eq 1
+        end
+
+        it 'can select hosts by properties on the host object using symbol as shorthand' do
+          expect(servers.roles_for(:app, select: :active).length).to eq 1
+        end
+
+        it 'can filter hosts by properties on the host using a regular proc' do
+          expect(servers.roles_for(:app, filter: lambda { |h| h.properties.active } ).length).to eq 1
+        end
+
+        it 'can select hosts by properties on the host using a regular proc' do
+          expect(servers.roles_for(:app, select: lambda { |h| h.properties.active } ).length).to eq 1
+        end
+
+        it 'raises if the regular proc filter would remove all matching hosts' do
+          I18n.expects(:t).with(:filter_removes_all_servers)
+          expect { servers.roles_for(:app, select: lambda { |h| h.properties.inactive }) }.to raise_error
+        end
+
+      end
     end
   end
 end
diff --git a/spec/lib/capistrano/configuration_spec.rb b/spec/lib/capistrano/configuration_spec.rb
index 514a0cde..724d9ff9 100644
--- a/spec/lib/capistrano/configuration_spec.rb
+++ b/spec/lib/capistrano/configuration_spec.rb
@@ -75,6 +75,8 @@ module Capistrano
         expect(config.fetch(:branch)).to eq question
       end
     end
-  end
 
+
+
+  end
 end
diff --git a/spec/lib/capistrano/dsl/env_spec.rb b/spec/lib/capistrano/dsl/env_spec.rb
index 002d1130..61414cba 100644
--- a/spec/lib/capistrano/dsl/env_spec.rb
+++ b/spec/lib/capistrano/dsl/env_spec.rb
@@ -5,79 +5,6 @@ module Capistrano
 
     describe Env do
 
-      let(:env) { Configuration.new }
-
-      describe '#role' do
-
-        it 'can add a role, with hosts' do
-          env.role(:app, %w{example.com})
-          env.roles_for(:app).first.hostname.should == "example.com"
-        end
-
-        it 'handles de-duplification within roles' do
-          env.role(:app, %w{example.com})
-          env.role(:app, %w{example.com})
-          env.roles_for(:app).length.should == 1
-        end
-
-        it 'accepts instances of server objects' do
-          pending
-          env.role(:app, [Capistrano::Configuration::Server.new('example.net'), 'example.com'])
-          env.roles_for(:app).length.should == 2
-        end
-
-        it 'accepts non-enumerable types' do
-          env.role(:app, 'example.com')
-          env.roles_for(:app).length.should == 1
-        end
-
-      end
-
-      describe '#server' do
-
-        it "can create a server with properties" do
-          env.server('example.com', roles: [:app, "web"], my: :value)
-          env.roles_for(:app).first.hostname.should == 'example.com'
-          env.roles_for(:web).first.hostname.should == 'example.com'
-          env.roles_for(:all).first.properties.my.should == :value
-        end
-
-      end
-
-      describe '#roles' do
-
-        before do
-          env.server('example.com', roles: :app, active: true)
-          env.server('example.org', roles: :app)
-        end
-
-        it 'raises if the filter would remove all matching hosts' do
-          pending
-          env.server('example.org', active: true)
-          lambda do
-            env.roles_for(:app, filter: lambda { |s| !s.properties.active })
-          end.should raise_error
-        end
-
-        it 'can filter hosts by properties on the host object using symbol as shorthand' do
-          env.roles_for(:app, filter: :active).length.should == 1
-        end
-
-        it 'can select hosts by properties on the host object using symbol as shorthand' do
-          env.roles_for(:app, select: :active).length.should == 1
-        end
-
-        it 'can filter hosts by properties on the host using a regular proc' do
-          env.roles_for(:app, filter: lambda { |h| h.properties.active } ).length.should == 1
-        end
-
-        it 'can select hosts by properties on the host using a regular proc' do
-          env.roles_for(:app, select: lambda { |h| h.properties.active } ).length.should == 1
-        end
-
-      end
-
     end
-
   end
 end