lib/beaker/hypervisor/ec2_helper.rb in beaker-1.21.0 vs lib/beaker/hypervisor/ec2_helper.rb in beaker-2.0.0

@@ -1,29 +1,41 @@
 module Beaker
   class EC2Helper
     # Return a list of open ports for testing based on a hosts role
     #
     # @todo horribly hard-coded
-    # @param [Array<String>] roles An array of roles
+    # @param [Host] host to find ports for
     # @return [Array<Number>] array of port numbers
     # @api private
-    def self.amiports(roles)
+    def self.amiports(host)
       ports = [22, 61613, 8139]
 
+      roles = host['roles']
+
       if roles.include? 'database'
+        ports << 5432
         ports << 8080
         ports << 8081
       end
 
       if roles.include? 'master'
         ports << 8140
       end
 
       if roles.include? 'dashboard'
         ports << 443
+        ports << 4433
         ports << 4435
       end
 
-      ports
+      # If they only specified one port in the host config file, YAML will have converted it
+      # into a string, but if it was more than one, an array.
+      user_ports = []
+      if host.has_key?('additional_ports')
+        user_ports = host['additional_ports'].is_a?(Array) ? host['additional_ports'] : [host['additional_ports']]
+      end
+
+      additional_ports = ports + user_ports
+      additional_ports.uniq
     end
   end
 end