lib/toxiproxy.rb in toxiproxy-0.1.3 vs lib/toxiproxy.rb in toxiproxy-0.1.4

@@ -1,15 +1,18 @@
 require "json"
 require "uri"
 require "net/http"
+require "forwardable"
 
 require "toxiproxy/collection"
 require "toxiproxy/toxic"
 require "toxiproxy/toxic_collection"
 
 class Toxiproxy
-  URI = ::URI.parse("http://127.0.0.1:8474")
+  extend SingleForwardable
+
+  DEFAULT_URI = 'http://127.0.0.1:8474'
   VALID_DIRECTIONS = [:upstream, :downstream]
 
   class NotFound < StandardError; end
   class ProxyExists < StandardError; end
   class InvalidToxic < StandardError; end
@@ -21,20 +24,11 @@
     @listen   = options[:listen] || "localhost:0"
     @name     = options[:name]
     @enabled  = options[:enabled]
   end
 
-  # Forwardable doesn't support delegating class methods, so we resort to
-  # `define_method` to delegate from Toxiproxy to #all, and from there to the
-  # proxy collection.
-  class << self
-    Collection::METHODS.each do |method|
-      define_method(method) do |*args, &block|
-        self.all.send(method, *args, &block)
-      end
-    end
-  end
+  def_delegators :all, *Collection::METHODS
 
   # Re-enables all proxies and disables all toxics.
   def self.reset
     request = Net::HTTP::Get.new("/reset")
     response = http.request(request)
@@ -58,10 +52,15 @@
     }
 
     Collection.new(proxies)
   end
 
+  # Sets the toxiproxy host to use.
+  def self.host=(host)
+    @uri = host.is_a?(::URI) ? host : ::URI.parse(host)
+  end
+
   # Convenience method to create a proxy.
   def self.create(options)
     self.new(options).create
   end
 
@@ -97,11 +96,11 @@
       self.create(proxy) unless existing
     }.compact
   end
 
   def self.running?
-    TCPSocket.new(URI.host, URI.port).close
+    TCPSocket.new(uri.host, uri.port).close
     true
   rescue Errno::ECONNREFUSED, Errno::ECONNRESET
     false
   end
 
@@ -124,15 +123,13 @@
   # Simulates the endpoint is down, by closing the connection and no
   # longer accepting connections. This is useful to simulate critical system
   # failure, such as a data store becoming completely unavailable.
   def down(&block)
     disable
-    begin
-      yield
-    ensure
-      enable
-    end
+    yield
+  ensure
+    enable
   end
 
   # Disables a Toxiproxy. This will drop all active connections and stop the proxy from listening.
   def disable
     request = Net::HTTP::Post.new("/proxies/#{name}")
@@ -205,11 +202,15 @@
     toxics
   end
 
   private
 
+  def self.uri
+    @uri ||= ::URI.parse(DEFAULT_URI)
+  end
+
   def self.http
-    @http ||= Net::HTTP.new(URI.host, URI.port)
+    @http ||= Net::HTTP.new(uri.host, uri.port)
   end
 
   def http
     self.class.http
   end