Greetings!

Please review the pull request #163: Ticket/master/3669 make puppet honor dns srv records opened by (jhelwig)

Some more information about the pull request:

  • Opened: Mon Oct 10 22:26:20 UTC 2011
  • Based on: puppetlabs:master (941e7945925f59e33dab2a09be505ea43ae8b3f2)
  • Requested merge: jhelwig:ticket/master/3669-make-puppet-honor-DNS-SRV-records (725b5a560ab13cd0ea09388d0cbe20a13730cec8)

Description:

This adds two new configuration variables:

  • use_srv_records: Will attempt to lookup SRV records for hostname
    found in srv_record (default: true)

  • srv_domain: The domain that will be queried for SRV records,
    (default: $domain)

If use_srv_records is set to true, then Puppet will attempt to find
the list of servers to use from SRV records on the domain specified
via srv_domain. The CA, report, and file servers can all be specified
via independent SRV records from the SRV records to use for looking up
the catalog server.

The SRV records must be for hosts in the form:

_x-puppet._tcp.$srv_domain
_x-puppet-ca._tcp.$srv_domain
_x-puppet-report._tcp.$srv_domain
_x-puppet-fileserver._tcp.$srv_domain

If no records are found for the _x-puppet-ca, _x-puppet-report, or
_x-puppet-fileserver services, then the SRV records for the _x-puppet
service will be used. However, if records exist for any of the more
specific services, Puppet will not attempt to use the _x-puppet
service to find an applicable server, even if none of the servers for
the more specific service can be contacted.

If Puppet is unable to connect to any of the servers specified in the
SRV records, then it will attempt to connect to the "normal" servers
settable via puppet.conf.

Signed-off-by: Jacob Helwig [email protected]
Signed-off-by: Daniel Pittman [email protected]
Signed-off-by: Andrew Forgue [email protected]

Thanks!
The Pull Request Bot

Diff follows:

diff --git a/lib/puppet/defaults.rb b/lib/puppet/defaults.rb
index 6d2ac86..5dcbd77 100644
--- a/lib/puppet/defaults.rb
+++ b/lib/puppet/defaults.rb
@@ -555,7 +555,12 @@ module Puppet
       :mode => 0640,
       :desc => "The log file for puppet agent.  This is generally not used."
     },
-    :server => ["puppet", "The server to which server puppet agent should connect"],
+    :server => {
+      :default => "puppet",
+      :desc => "The server to which the puppet agent should connect"
+    },
+    :use_srv_records => [true, "Whether the server will search for SRV records in DNS for the current domain."],
+    :srv_domain => [ "#{domain}", "The domain which will be queried to find the SRV records of servers to use."],
     :ignoreschedules => [false,
       "Boolean; whether puppet agent should ignore schedules.  This is useful
       for initial puppet agent runs."],
@@ -568,7 +573,7 @@ module Puppet
       it with the `--no-client` option."],
     :listen => [false, "Whether puppet agent should listen for
       connections.  If this is true, then puppet agent will accept incoming
-      REST API requests, subject to the default ACLs and the ACLs set in 
+      REST API requests, subject to the default ACLs and the ACLs set in
       the `rest_authconfig` file. Puppet agent can respond usefully to
       requests on the `run`, `facts`, `certificate`, and `resource` endpoints."],
     :ca_server => ["$server", "The server to use for certificate
diff --git a/lib/puppet/indirector/certificate/rest.rb b/lib/puppet/indirector/certificate/rest.rb
index 921b857..1cbf06a 100644
--- a/lib/puppet/indirector/certificate/rest.rb
+++ b/lib/puppet/indirector/certificate/rest.rb
@@ -6,6 +6,7 @@ class Puppet::SSL::Certificate::Rest < Puppet::Indirector::REST
 
   use_server_setting(:ca_server)
   use_port_setting(:ca_port)
+  use_srv_service(:ca)
 
   def find(request)
     return nil unless result = super
diff --git a/lib/puppet/indirector/file_content/rest.rb b/lib/puppet/indirector/file_content/rest.rb
index 59975d3..1f6bc6e 100644
--- a/lib/puppet/indirector/file_content/rest.rb
+++ b/lib/puppet/indirector/file_content/rest.rb
@@ -4,4 +4,6 @@ require 'puppet/indirector/rest'
 
 class Puppet::Indirector::FileContent::Rest < Puppet::Indirector::REST
   desc "Retrieve file contents via a REST HTTP interface."
+
+  use_srv_service(:fileserver)
 end
diff --git a/lib/puppet/indirector/report/rest.rb b/lib/puppet/indirector/report/rest.rb
index 601da9e..a40a9f8 100644
--- a/lib/puppet/indirector/report/rest.rb
+++ b/lib/puppet/indirector/report/rest.rb
@@ -4,4 +4,5 @@ class Puppet::Transaction::Report::Rest < Puppet::Indirector::REST
   desc "Get server report over HTTP via REST."
   use_server_setting(:report_server)
   use_port_setting(:report_port)
+  use_srv_service(:report)
 end
diff --git a/lib/puppet/indirector/rest.rb b/lib/puppet/indirector/rest.rb
index 19daff5..e02eddb 100644
--- a/lib/puppet/indirector/rest.rb
+++ b/lib/puppet/indirector/rest.rb
@@ -4,6 +4,7 @@ require 'uri'
 require 'puppet/network/http_pool'
 require 'puppet/network/http/api/v1'
 require 'puppet/network/http/compression'
+require 'puppet/network/resolver'
 
 # Access objects via REST
 class Puppet::Indirector::REST < Puppet::Indirector::Terminus
@@ -19,19 +20,56 @@ class Puppet::Indirector::REST < Puppet::Indirector::Terminus
     @server_setting = setting
   end
 
-  def self.server
-    Puppet.settings[server_setting || :server]
-  end
-
   # Specify the setting that we should use to get the port.
   def self.use_port_setting(setting)
     @port_setting = setting
   end
 
+  # Specify the service to use when doing SRV record lookup
+  def self.use_srv_service(service)
+    @srv_service = service
+  end
+
+  def self.srv_service
+    @srv_service || :puppet
+  end
+
+  def self.server
+    Puppet.settings[server_setting || :server]
+  end
+
   def self.port
     Puppet.settings[port_setting || :masterport].to_i
   end
 
+  def resolve_servers_for(request)
+    # We were given a specific server to use, so just use that one.
+    # This happens if someone does something like specifying a file
+    # source using a puppet:// URI with a specific server.
+    return yield request if !request.server.nil?
+
+    if Puppet.settings[:use_srv_records]
+      Puppet::Network::Resolver.each_srv_record(
+        Puppet.settings[:srv_domain],
+        self.class.srv_service
+      ) do |srv_server, srv_port|
+        begin
+          request.server = srv_server
+          request.port   = srv_port
+          return yield request
+        rescue SystemCallError => e
+          Puppet.warning "Error connecting to #{srv_server}:#{srv_port}: #{e.message}"
+        end
+      end
+    end
+
+    # ... Fall back onto the default server.
+    Puppet.debug "No more servers left, falling back to #{self.class.server}" if Puppet.settings[:use_srv_records]
+    request.server = self.class.server
+    request.port = self.class.port
+    return yield request
+  end
+
   # Figure out the content type, turn that into a format, and use the format
   # to extract the body of the response.
   def deserialize(response, multiple = false)
@@ -110,20 +148,29 @@ class Puppet::Indirector::REST < Puppet::Indirector::Terminus
   def find(request)
     uri, body = request_to_uri_and_body(request)
     uri_with_query_string = "#{uri}?#{body}"
-    # WEBrick in Ruby 1.9.1 only supports up to 1024 character lines in an HTTP request
-    # http://redmine.ruby-lang.org/issues/show/3991
-    response = if "GET #{uri_with_query_string} HTTP/1.1\r\n".length > 1024
-      http_post(request, uri, body, headers)
-    else
-      http_get(request, uri_with_query_string, headers)
+
+    response = resolve_servers_for(request) do |request|
+      # WEBrick in Ruby 1.9.1 only supports up to 1024 character lines in an HTTP request
+      # http://redmine.ruby-lang.org/issues/show/3991
+      if "GET #{uri_with_query_string} HTTP/1.1\r\n".length > 1024
+        http_post(request, uri, body, headers)
+      else
+        http_get(request, uri_with_query_string, headers)
+      end
     end
     result = deserialize response
+
+    return nil unless result
+
     result.name = request.key if result.respond_to?(:name=)
     result
   end
 
   def head(request)
-    response = http_head(request, indirection2uri(request), headers)
+    response = resolve_servers_for(request) do |request|
+      http_head(request, indirection2uri(request), headers)
+    end
+
     case response.code
     when "404"
       return false
@@ -136,20 +183,28 @@ class Puppet::Indirector::REST < Puppet::Indirector::Terminus
   end
 
   def search(request)
-    unless result = deserialize(http_get(request, indirection2uri(request), headers), true)
-      return []
+    result = resolve_servers_for(request) do |request|
+      deserialize(http_get(request, indirection2uri(request), headers), true)
     end
-    result
+
+    # result from the server can be nil, but we promise to return an array...
+    result || []
   end
 
   def destroy(request)
     raise ArgumentError, "DELETE does not accept options" unless request.options.empty?
-    deserialize http_delete(request, indirection2uri(request), headers)
+
+    resolve_servers_for(request) do |request|
+      return deserialize(http_delete(request, indirection2uri(request), headers))
+    end
   end
 
   def save(request)
     raise ArgumentError, "PUT does not accept options" unless request.options.empty?
-    deserialize http_put(request, indirection2uri(request), request.instance.render, headers.merge({ "Content-Type" => request.instance.mime }))
+
+    resolve_servers_for(request) do |request|
+      deserialize http_put(request, indirection2uri(request), request.instance.render, headers.merge({ "Content-Type" => request.instance.mime }))
+    end
   end
 
   private
diff --git a/lib/puppet/network/resolver.rb b/lib/puppet/network/resolver.rb
new file mode 100644
index 0000000..918b694
--- /dev/null
+++ b/lib/puppet/network/resolver.rb
@@ -0,0 +1,80 @@
+require 'resolv'
+module Puppet::Network; end
+
+module Puppet::Network::Resolver
+  # Iterate through the list of servers that service this hostname
+  # and yield each server/port since SRV records have ports in them
+  # It will override whatever masterport setting is already set.
+  def self.each_srv_record(domain, service_name = :puppet, &block)
+    Puppet.debug "Searching for SRV records for domain: #{domain}"
+
+    case service_name
+      when :puppet then service = '_x-puppet'
+      when :ca     then service = '_x-puppet-ca'
+      when :report then service = '_x-puppet-report'
+      when :file   then service = '_x-puppet-fileserver'
+    end
+    srv_record = "#{service}._tcp.#{domain}"
+
+    resolver = Resolv::DNS.new
+    records = resolver.getresources(srv_record, Resolv::DNS::Resource::IN::SRV)
+    Puppet.debug "Found #{records.size} SRV records for: #{srv_record}"
+
+    if records.size == 0 && service_name != :puppet
+      # Try the generic :puppet service if no SRV records were found
+      # for the specific service.
+      each_srv_record(domain, :puppet, &block)
+    else
+      each_priority(records) do |priority, records|
+        while next_rr = records.delete(find_weighted_server(records))
+          Puppet.debug "Yielding next server of #{next_rr.target.to_s}:#{next_rr.port}"
+          yield next_rr.target.to_s, next_rr.port
+        end
+      end
+    end
+  end
+
+  private
+
+  def self.each_priority(records)
+    pri_hash = records.inject({}) do |groups, element|
+      groups[element.priority] ||= []
+      groups[element.priority] << element
+      groups
+    end
+
+    pri_hash.keys.sort.each do |key|
+      yield key, pri_hash[key]
+    end
+  end
+
+  def self.find_weighted_server(records)
+    return nil if records.nil? || records.empty?
+    return records.first if records.size == 1
+
+    # Calculate the sum of all weights in the list of resource records,
+    # This is used to then select hosts until the weight exceeds what
+    # random number we selected.  For example, if we have weights of 1 8 and 3:
+    #
+    # |-|---|--------|
+    #        ^
+    # We generate a random number 5, and iterate through the records, adding
+    # the current record's weight to the accumulator until the weight of the
+    # current record plus previous records is greater than the random number.
+
+    total_weight = records.inject(0) { |sum,record|
+      sum + weight(record)
+    }
+    current_weight = 0
+    chosen_weight  = 1 + Kernel.rand(total_weight)
+
+    records.each do |record|
+      current_weight += weight(record)
+      return record if current_weight >= chosen_weight
+    end
+  end
+
+  def self.weight(record)
+    record.weight == 0 ? 1 : record.weight * 10
+  end
+end
diff --git a/lib/puppet/type/file/content.rb b/lib/puppet/type/file/content.rb
index 93b8e69..da51ff4 100755
--- a/lib/puppet/type/file/content.rb
+++ b/lib/puppet/type/file/content.rb
@@ -199,10 +199,29 @@ module Puppet
       end
     end
 
-    def chunk_file_from_source(source_or_content)
+
+    def get_from_source(source_or_content)
       request = Puppet::Indirector::Request.new(:file_content, :find, source_or_content.full_path.sub(/^\//,''))
+
+      if Puppet.settings[:use_srv_records] and ! source_or_content.server? then
+        Puppet::Network::Resolver.each_srv_record(Puppet.settings[:srv_record]) do |server, port|
+          begin
+            connection = Puppet::Network::HttpPool.http_instance(server, port)
+            return yield connection.request_get(indirection2uri(request), add_accept_encoding({"Accept" => "raw"}))
+          rescue SystemCallError => e
+            Puppet.warning "Error connecting to #{server}:#{port}: #{e.message}"
+          end
+        end
+        Puppet.debug "No more servers left, falling back to #{source_or_content.server}"
+      end
+
       connection = Puppet::Network::HttpPool.http_instance(source_or_content.server, source_or_content.port)
-      connection.request_get(indirection2uri(request), add_accept_encoding({"Accept" => "raw"})) do |response|
+      return yield connection.request_get(indirection2uri(request), add_accept_encoding({"Accept" => "raw"}))
+    end
+
+
+    def chunk_file_from_source(source_or_content)
+      get_from_source(source_or_content) do |response|
         case response.code
         when /^2/;  uncompress(response) { |uncompressor| response.read_body { |chunk| yield uncompressor.uncompress(chunk) } }
         else
diff --git a/lib/puppet/type/file/source.rb b/lib/puppet/type/file/source.rb
index 2080b9e..72590d1 100755
--- a/lib/puppet/type/file/source.rb
+++ b/lib/puppet/type/file/source.rb
@@ -183,6 +183,10 @@ module Puppet
       Puppet::Util.uri_to_path(uri) if found?
     end
 
+    def server?
+       uri and uri.host
+    end
+
     def server
       (uri and uri.host) or Puppet.settings[:server]
     end
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index e0227be..aec998e 100755
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -100,7 +100,7 @@ RSpec.configure do |config|
         indirector.instance_variable_set(variable, value)
       end
     end
-    $saved_indirection_state = nil
+    $saved_indirection_state = {}
 
     GC.enable
   end
diff --git a/spec/unit/indirector/certificate/rest_spec.rb b/spec/unit/indirector/certificate/rest_spec.rb
index 870d9e4..4ec475b 100755
--- a/spec/unit/indirector/certificate/rest_spec.rb
+++ b/spec/unit/indirector/certificate/rest_spec.rb
@@ -20,6 +20,10 @@ describe Puppet::SSL::Certificate::Rest do
     Puppet::SSL::Certificate::Rest.port_setting.should == :ca_port
   end
 
+  it "should use the :ca SRV service" do
+    Puppet::SSL::Certificate::Rest.srv_service.should == :ca
+  end
+
   it "should make sure found certificates have their names set to the search string" do
     terminus = Puppet::SSL::Certificate::Rest.new
 
diff --git a/spec/unit/indirector/file_content/rest_spec.rb b/spec/unit/indirector/file_content/rest_spec.rb
index 06ad16e..99925a8 100755
--- a/spec/unit/indirector/file_content/rest_spec.rb
+++ b/spec/unit/indirector/file_content/rest_spec.rb
@@ -1,10 +1,14 @@
 #!/usr/bin/env rspec
 require 'spec_helper'
 
-require 'puppet/indirector/file_content'
+require 'puppet/indirector/file_content/rest'
 
-describe "Puppet::Indirector::Content::Rest" do
+describe Puppet::Indirector::FileContent::Rest do
   it "should add the node's cert name to the arguments"
 
   it "should set the content type to text/plain"
+
+  it "should use the :fileserver SRV service" do
+    Puppet::Indirector::FileContent::Rest.srv_service.should == :fileserver
+  end
 end
diff --git a/spec/unit/indirector/report/rest_spec.rb b/spec/unit/indirector/report/rest_spec.rb
index 806ae6e..8e969ab 100755
--- a/spec/unit/indirector/report/rest_spec.rb
+++ b/spec/unit/indirector/report/rest_spec.rb
@@ -24,4 +24,8 @@ describe Puppet::Transaction::Report::Rest do
     Puppet::Transaction::Report::Rest.server.should_not be_nil
     Puppet::Transaction::Report::Rest.port.should_not be_nil
   end
+
+  it "should use the :report SRV service" do
+    Puppet::Transaction::Report::Rest.srv_service.should == :report
+  end
 end
diff --git a/spec/unit/indirector/rest_spec.rb b/spec/unit/indirector/rest_spec.rb
index 042b7ca..34ac84d 100755
--- a/spec/unit/indirector/rest_spec.rb
+++ b/spec/unit/indirector/rest_spec.rb
@@ -138,6 +138,10 @@ describe Puppet::Indirector::REST do
     end
   end
 
+  it 'should default to :puppet for the srv_service' do
+    Puppet::Indirector::REST.srv_service.should == :puppet
+  end
+
   describe "when deserializing responses" do
     it "should return nil if the response code is 404" do
       response = mock 'response'
@@ -265,6 +269,109 @@ describe Puppet::Indirector::REST do
     end
   end
 
+  describe "when resolving servers for requests" do
+    before :each do
+      @request = Puppet::Indirector::Request.new(:indirection, :method, :key)
+      @block_return = 'returned from the block'
+    end
+
+    it "should return the request unmodified when not using SRV records" do
+      @request.server = 'puppet.example.com'
+      @request.port   = '90210'
+
+      count = 0
+      rval = @searcher.resolve_servers_for(@request) do |got|
+        count += 1
+        got.server.should == @request.server
+        got.port.should   == @request.port
+        @block_return
+      end
+      count.should == 1
+
+      rval.should == @block_return
+    end
+
+    it "should tell the resolver which service to use to look up SRV records" do
+      @rest_class.use_srv_service(:report)
+      Puppet::Network::Resolver.expects(:each_srv_record).with(
+        Puppet.settings[:srv_domain],
+        :report
+      ).returns(nil)
+
+      @searcher.resolve_servers_for(@request) {|x| }
+    end
+
+    it "should default to using the :puppet service to look up SRV records" do
+      @rest_class.use_srv_service(nil)
+      Puppet::Network::Resolver.expects(:each_srv_record).with(
+        Puppet.settings[:srv_domain],
+        :puppet
+      ).returns(nil)
+
+      @searcher.resolve_servers_for(@request) {|x| }
+    end
+
+    it "uses the specified server if the request already has one" do
+      @request.server = 'foo'
+      @searcher.resolve_servers_for(@request) do |req|
+        req.server.should == 'foo'
+      end
+    end
+
+    describe "when SRV returns servers" do
+      before :each do
+        @dns_mock = mock('dns')
+        Resolv::DNS.expects(:new).returns(@dns_mock)
+
+        @port = 7205
+        @host = '_x-puppet._tcp.example.com'
+        @srv_records = [Resolv::DNS::Resource::IN::SRV.new(0, 0, @port, @host)]
+
+        Puppet.settings[:use_srv_records] = true
+        @dns_mock.expects(:getresources).
+          with("_x-puppet._tcp.#{Puppet.settings[:srv_domain]}", Resolv::DNS::Resource::IN::SRV).
+          returns(@srv_records)
+      end
+
+      it "should return the SRV record when found" do
+        count = 0
+        rval = @searcher.resolve_servers_for(@request) do |got|
+          count += 1
+          got.server.should == '_x-puppet._tcp.example.com'
+          got.port.should == 7205
+
+          @block_return
+        end
+        count.should == 1
+
+        rval.should == @block_return
+      end
+
+      it "should fall back to the default server when the block raises a SystemCallError" do
+        count = 0
+        second_pass = nil
+
+        rval = @searcher.resolve_servers_for(@request) do |got|
+          count += 1
+
+          if got.server == '_x-puppet._tcp.example.com' then
+            raise SystemCallError, "example failure"
+          else
+            second_pass = got
+          end
+
+          @block_return
+        end
+
+        second_pass.server.should == 'puppet'
+        second_pass.port.should   == 8140
+        count.should == 2
+
+        rval.should == @block_return
+      end
+    end
+  end
+
   describe "when doing a find" do
     before :each do
       @connection = stub('mock http connection', :get => @response, :verify_callback= => nil)
diff --git a/spec/unit/network/resolver_spec.rb b/spec/unit/network/resolver_spec.rb
new file mode 100755
index 0000000..434c801
--- /dev/null
+++ b/spec/unit/network/resolver_spec.rb
@@ -0,0 +1,185 @@
+#!/usr/bin/env ruby
+require File.dirname(__FILE__) + '/../../spec_helper'
+require 'puppet/network/resolver'
+
+describe Puppet::Network::Resolver do
+  before do
+    @dns_mock_object = mock('dns')
+    Resolv::DNS.stubs(:new).returns(@dns_mock_object)
+
+    @rr_type         = Resolv::DNS::Resource::IN::SRV
+    @test_srv_domain = "domain.com"
+    @test_a_hostname = "puppet.domain.com"
+    @test_port       = 1000
+
+    # The records we should use.
+    @test_records = [
+      #                                  priority,  weight, port, hostname
+      Resolv::DNS::Resource::IN::SRV.new(0,         20,     8140, "puppet1.domain.com"),
+      Resolv::DNS::Resource::IN::SRV.new(0,         80,     8140, "puppet2.domain.com"),
+      Resolv::DNS::Resource::IN::SRV.new(1,         1,      8140, "puppet3.domain.com"),
+      Resolv::DNS::Resource::IN::SRV.new(4,         1,      8140, "puppet4.domain.com")
+    ]
+  end
+
+  describe "when resolving a host without SRV records" do
+    it "should not yield anything" do
+      # No records returned for a DNS entry without any SRV records
+      @dns_mock_object.expects(:getresources).with(
+        "_x-puppet._tcp.#{@test_a_hostname}",
+        @rr_type
+      ).returns([])
+
+      Puppet::Network::Resolver.each_srv_record(@test_a_hostname) do |hostname, port, remaining|
+        fail_with "host with no records passed block"
+      end
+    end
+  end
+
+  describe "when resolving a host with SRV records" do
+    it "should iterate through records in priority order" do
+      # The order of the records that should be returned,
+      # an array means unordered (for weight)
+      order = {
+        0 => ["puppet1.domain.com", "puppet2.domain.com"],
+        1 => ["puppet3.domain.com"],
+        2 => ["puppet4.domain.com"]
+      }
+
+      @dns_mock_object.expects(:getresources).with(
+        "_x-puppet._tcp.#{@test_srv_domain}",
+        @rr_type
+      ).returns(@test_records)
+
+      Puppet::Network::Resolver.each_srv_record(@test_srv_domain) do |hostname, port|
+        expected_priority = order.keys.min
+
+        order[expected_priority].should include(hostname)
+        port.should_not be(@test_port)
+
+        # Remove the host from our expected hosts
+        order[expected_priority].delete hostname
+
+        # Remove this priority level if we're done with it
+        order.delete expected_priority if order[expected_priority] == []
+      end
+    end
+
+    it "should fall back to the :puppet service if no records are found for a more specific service" do
+      # The order of the records that should be returned,
+      # an array means unordered (for weight)
+      order = {
+        0 => ["puppet1.domain.com", "puppet2.domain.com"],
+        1 => ["puppet3.domain.com"],
+        2 => ["puppet4.domain.com"]
+      }
+
+      @dns_mock_object.expects(:getresources).with(
+        "_x-puppet-report._tcp.#{@test_srv_domain}",
+        @rr_type
+      ).returns([])
+
+      @dns_mock_object.expects(:getresources).with(
+        "_x-puppet._tcp.#{@test_srv_domain}",
+        @rr_type
+      ).returns(@test_records)
+
+      Puppet::Network::Resolver.each_srv_record(@test_srv_domain, :report) do |hostname, port|
+        expected_priority = order.keys.min
+
+        order[expected_priority].should include(hostname)
+        port.should_not be(@test_port)
+
+        # Remove the host from our expected hosts
+        order[expected_priority].delete hostname
+
+        # Remove this priority level if we're done with it
+        order.delete expected_priority if order[expected_priority] == []
+      end
+    end
+
+    it "should use SRV records from the specific service if they exist" do
+      # The order of the records that should be returned,
+      # an array means unordered (for weight)
+      order = {
+        0 => ["puppet1.domain.com", "puppet2.domain.com"],
+        1 => ["puppet3.domain.com"],
+        2 => ["puppet4.domain.com"]
+      }
+
+      bad_records = [
+        #                                  priority,  weight, port, hostname
+        Resolv::DNS::Resource::IN::SRV.new(0,         20,     8140, "puppet1.bad.domain.com"),
+        Resolv::DNS::Resource::IN::SRV.new(0,         80,     8140, "puppet2.bad.domain.com"),
+        Resolv::DNS::Resource::IN::SRV.new(1,         1,      8140, "puppet3.bad.domain.com"),
+        Resolv::DNS::Resource::IN::SRV.new(4,         1,      8140, "puppet4.bad.domain.com")
+      ]
+
+      @dns_mock_object.expects(:getresources).with(
+        "_x-puppet-report._tcp.#{@test_srv_domain}",
+        @rr_type
+      ).returns(@test_records)
+
+      @dns_mock_object.stubs(:getresources).with(
+        "_x-puppet._tcp.#{@test_srv_domain}",
+        @rr_type
+      ).returns(bad_records)
+
+      Puppet::Network::Resolver.each_srv_record(@test_srv_domain, :report) do |hostname, port|
+        expected_priority = order.keys.min
+
+        order[expected_priority].should include(hostname)
+        port.should_not be(@test_port)
+
+        # Remove the host from our expected hosts
+        order[expected_priority].delete hostname
+
+        # Remove this priority level if we're done with it
+        order.delete expected_priority if order[expected_priority] == []
+      end
+    end
+  end
+
+  describe "when finding weighted servers" do
+    it "should return nil when no records were found" do
+      Puppet::Network::Resolver.find_weighted_server([]).should == nil
+    end
+
+    it "should return the first record when one record is passed" do
+      result = Puppet::Network::Resolver.find_weighted_server([@test_records.first])
+      result.should == @test_records.first
+    end
+
+    {
+      "all have weights"  => [1, 3, 2, 4],
+      "some have weights" => [2, 0, 1, 0],
+      "none have weights" => [0, 0, 0, 0],
+    }.each do |name, weights|
+      it "should return correct results when #{name}" do
+        records = []
+        count   = 0
+        weights.each do |w|
+          count += 1
+          #                                             priority, weight, port, server
+          records << Resolv::DNS::Resource::IN::SRV.new(0,        w,      1,    count.to_s)
+        end
+
+        seen  = Hash.new(0)
+        total_weight = records.inject(0) { |sum, record|
+          sum + Puppet::Network::Resolver.weight(record)
+        }
+
+        total_weight.times do |n|
+          Kernel.expects(:rand).once.with(total_weight).returns(n)
+          server = Puppet::Network::Resolver.find_weighted_server(records)
+          seen[server] += 1
+        end
+
+        seen.length.should == records.length
+        records.each do |record|
+          seen[record].should == Puppet::Network::Resolver.weight(record)
+        end
+      end
+    end
+  end
+end
diff --git a/spec/unit/type/file/content_spec.rb b/spec/unit/type/file/content_spec.rb
index b208240..df9ca21 100755
--- a/spec/unit/type/file/content_spec.rb
+++ b/spec/unit/type/file/content_spec.rb
@@ -1,6 +1,8 @@
 #!/usr/bin/env rspec
 require 'spec_helper'
 
+require 'puppet/network/resolver'
+
 content = Puppet::Type.type(:file).attrclass(:content)
 describe content do
   include PuppetSpec::Files
@@ -331,7 +333,7 @@ describe content do
         @response.stubs(:read_body).multiple_yields(*(["source file content"]*10000))
 
         @conn = stub_everything 'connection'
-        @conn.stubs(:request_get).yields(@response)
+        @conn.stubs(:request_get).returns(@response)
         Puppet::Network::HttpPool.stubs(:http_instance).returns @conn
 
         @content = @resource.newattr(:content)
diff --git a/test/lib/puppettest.rb b/test/lib/puppettest.rb
index 6bae80a..4d32f14 100755
--- a/test/lib/puppettest.rb
+++ b/test/lib/puppettest.rb
@@ -154,9 +154,8 @@ module PuppetTest
     end
 
 
-      @configpath = File.join(
-        tmpdir,
-
+    @configpath = File.join(
+      tmpdir,
       "configdir" + @@testcount.to_s + "/"
     )
 
@@ -233,11 +232,11 @@ module PuppetTest
   def tmpdir
     unless @tmpdir
       @tmpdir = case Facter["operatingsystem"].value
-        when "Darwin"; "/private/tmp"
-        when "SunOS"; "/var/tmp"
+        when "Darwin" then "/private/tmp"
+        when "SunOS"  then "/var/tmp"
         else
           "/tmp"
-            end
+      end
 
 
       @tmpdir = File.join(@tmpdir, "puppettesting#{Process.pid}")

    

--
You received this message because you are subscribed to the Google Groups "Puppet Developers" group.
To post to this group, send email to [email protected].
To unsubscribe from this group, send email to [email protected].
For more options, visit this group at http://groups.google.com/group/puppet-dev?hl=en.

Reply via email to