Add CC::Resolv

wfleming · wfleming · commit 2330cf2e0cf4 · 2016-11-10T18:37:52.000-05:00
diff --git a/lib/cc/resolv.rb b/lib/cc/resolv.rb
@@ -0,0 +1,39 @@
+require "resolv-replace"
+
+module CC
+  class Resolv
+    def self.with_fixed_dns(dns = ::Resolv::DNS.new)
+      ::Resolv::DefaultResolver.replace_resolvers([Fixed.new(dns)])
+
+      yield if block_given?
+    ensure
+      # There's no way to ask what the current values are before we override
+      # them; hopefully going by the source is good enough.
+      # https://docs.ruby-lang.org/en/2.0.0/Resolv.html#method-c-new
+      default_resolvers = [::Resolv::Hosts.new, ::Resolv::DNS.new]
+      ::Resolv::DefaultResolver.replace_resolvers(default_resolvers)
+    end
+
+    class Fixed
+      def initialize(fallback)
+        @addresses = {}
+        @fallback = fallback
+      end
+
+      def each_address(name)
+        if addresses.key?(name)
+          yield addresses.fetch(name)
+        else
+          fallback.each_address(name) do |address|
+            addresses[name] ||= address
+            yield address
+          end
+        end
+      end
+
+      private
+
+      attr_reader :addresses, :fallback
+    end
+  end
+end
diff --git a/spec/cc/resolv_spec.rb b/spec/cc/resolv_spec.rb
@@ -0,0 +1,44 @@
+require "spec_helper"
+require "cc/resolv"
+
+module CC
+  describe Resolv do
+    describe ".with_fixed_dns" do
+      it "replaces the default resolver for the duration of the block" do
+        fallback = double
+
+        expect(fallback).to receive(:each_address).
+          with("google.com").and_yield("overridden")
+
+        Resolv.with_fixed_dns(fallback) do
+          expect(::Resolv.getaddress("google.com")).to eq "overridden"
+          expect(::Resolv.getaddress("google.com")).to eq "overridden"
+        end
+
+        expect(::Resolv.getaddress("google.com")).not_to eq "overridden"
+      end
+    end
+
+    describe Resolv::Fixed do
+      describe "#each_address" do
+        it "delegates to the fallback resolver and caches the first address" do
+          fallback = double
+          fixed = Resolv::Fixed.new(fallback)
+
+          allow(fallback).to receive(:each_address).
+            with("host").once.
+            and_yield("address-1").
+            and_yield("address-2")
+
+          yielded_1 = []
+          yielded_2 = []
+          fixed.each_address("host") { |a| yielded_1 << a }
+          fixed.each_address("host") { |a| yielded_2 << a }
+
+          expect(yielded_1).to eq ["address-1", "address-2"]
+          expect(yielded_2).to eq ["address-1"]
+        end
+      end
+    end
+  end
+end
diff --git a/spec/support/resolv_helpers.rb b/spec/support/resolv_helpers.rb
@@ -0,0 +1,12 @@
+module ResolvHelpers
+  def stub_resolv(name, address)
+    dns = double
+    allow(::Resolv::DNS).to receive(:new).and_return(dns)
+    allow(dns).to receive(:each_address).
+      with(name).and_yield(Resolv::IPv4.create(address))
+  end
+end
+
+RSpec.configure do |conf|
+  conf.include(ResolvHelpers)
+end
