lib/infrataster/server.rb in infrataster-0.2.1 vs lib/infrataster/server.rb in infrataster-0.2.2

@@ -1,18 +1,27 @@
 require 'tmpdir'
 require 'net/ssh'
 require 'net/ssh/gateway'
 require 'ipaddr'
 require 'shellwords'
+require 'ostruct'
 
 module Infrataster
   class Server
     Error = Class.new(StandardError)
 
     class << self
 
-      def define(*args)
-        @@servers << Server.new(*args)
+      def define(name, *args, &block)
+        address = args.shift
+        options = args.any? ? args.shift : {}
+        if block
+          st = OpenStruct.new
+          block.call(st)
+          address = st.address if st.address
+          st.each_pair { |k, v| options[k] = v unless k == :address }
+        end
+        @@servers << Server.new(name, address, options)
       end
 
       def defined_servers
         @@servers
       end