wsdirector-core 1.0.1 → 1.0.2

data/lib/.rbnext/2.7/wsdirector/scenario_reader.rb

@@ -0,0 +1,172 @@
+# frozen_string_literal: true
+
+require "erb"
+require "json"
+require "wsdirector/ext/deep_dup"
+
+module WSDirector
+  # Read and parse different scenarios
+  class ScenarioReader
+    using WSDirector::Ext::DeepDup
+
+    include WSDirector::Utils
+
+    class << self
+      def parse(scenario, **options)
+        new(**options).parse(scenario)
+      end
+    end
+
+    attr_reader :default_connections_options, :locals
+
+    def initialize(scale: 1, connection_options: {}, locals: {})
+      @scale = scale
+      @default_connections_options = connection_options
+      @locals = locals
+    end
+
+    def parse(scenario)
+      contents =
+        if scenario.is_a?(String)
+          if File.file?(scenario)
+            parse_file(scenario)
+          else
+            parse_from_str(scenario).then do |_1|
+              next _1 if _1.is_a?(Array)
+              [_1]
+            end
+          end.flatten
+        else
+          scenario
+        end
+
+      contents.map! do |item|
+        item.is_a?(String) ? {item => {}} : item
+      end
+
+      if contents.first&.key?("client")
+        contents = transform_with_loop(contents, multiple: true)
+        parse_multiple_scenarios(contents)
+      else
+        contents = transform_with_loop(contents)
+        {"total" => 1, "clients" => [parse_simple_scenario(contents)]}
+      end
+    end
+
+    private
+
+    JSON_FILE_FORMAT = /.+.(json)\z/
+    private_constant :JSON_FILE_FORMAT
+
+    def parse_from_str(contents)
+      JSON.parse(contents)
+    rescue JSON::ParserError
+      parse_yaml(contents)
+    end
+
+    def parse_file(file)
+      if file.match?(JSON_FILE_FORMAT)
+        JSON.parse(File.read(file))
+      else
+        parse_yaml(file)
+      end
+    end
+
+    def parse_yaml(path)
+      contents = File.file?(path) ? File.read(path) : path
+
+      if defined?(ERB)
+        contents = ERB.new(contents).result(erb_context)
+      end
+
+      ::YAML.load(contents, aliases: true, permitted_classes: [Date, Time, Regexp]) || {}
+    rescue ArgumentError
+      ::YAML.load(contents) || {}
+    end
+
+    def handle_steps(steps)
+      steps.flat_map.with_index do |step, id|
+        if step.is_a?(Hash)
+          type, data = step.to_a.first
+
+          data["sample"] = [1, parse_multiplier(data["sample"])].max if data["sample"]
+
+          multiplier = parse_multiplier(data.delete("multiplier") || "1")
+          Array.new(multiplier) { {"type" => type, "id" => id}.merge(data) }
+        else
+          {"type" => step, "id" => id}
+        end
+      end
+    end
+
+    def parse_simple_scenario(
+      steps,
+      multiplier: 1, name: "default", ignore: nil, protocol: "base",
+      connection_options: {}
+    )
+      {
+        "multiplier" => multiplier,
+        "steps" => handle_steps(steps),
+        "name" => name,
+        "ignore" => ignore,
+        "protocol" => protocol,
+        "connection_options" => default_connections_options.merge(connection_options)
+      }
+    end
+
+    def parse_multiple_scenarios(definitions)
+      total_count = 0
+      clients = definitions.map.with_index do |client, i|
+        _, client = client.to_a.first
+        multiplier = parse_multiplier(client.delete("multiplier") || "1")
+        name = client.delete("name") || (i + 1).to_s
+        connection_options = client.delete("connection_options") || {}
+        total_count += multiplier
+        ignore = parse_ignore(client.fetch("ignore", nil))
+        protocol = client.fetch("protocol", "base")
+
+        parse_simple_scenario(
+          client.fetch("actions", []),
+          multiplier: multiplier,
+          name: name,
+          ignore: ignore,
+          protocol: protocol,
+          connection_options: connection_options
+        )
+      end
+      {"total" => total_count, "clients" => clients}
+    end
+
+    def transform_with_loop(contents, multiple: false)
+      contents.flat_map do |content|
+        loop_data = content.dig("client", "loop") || content.dig("loop")
+        next content unless loop_data
+
+        loop_multiplier = parse_multiplier(loop_data["multiplier"] || "1")
+
+        if multiple
+          content["client"]["actions"] = (loop_data["actions"] * loop_multiplier).map(&:deep_dup)
+          content
+        else
+          loop_data["actions"] * loop_multiplier
+        end
+      end
+    end
+
+    def parse_ignore(str)
+      return unless str
+
+      Array(str)
+    end
+
+    def erb_context
+      binding.then do |b|
+        locals.each do |_1, _2|
+          b.local_variable_set(_1, _2)
+        end
+
+        b
+      end
+    end
+  end
+end

data/lib/.rbnext/3.0/wsdirector/ext/formatting.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module WSDirector
+  module Ext
+    # Extend Object through refinements
+    module Formatting
+      refine ::Object do
+        def truncate(*) ;  itself; end
+      end
+
+      refine ::String do
+        def truncate(limit)
+          return self if size <= limit
+
+          "#{self[0..(limit - 3)]}..."
+        end
+      end
+
+      refine ::Hash do
+        def truncate(limit)
+          str = to_json
+
+          str.truncate(limit)
+        end
+      end
+
+      refine ::Float do
+        def duration
+          if self > 1
+            "#{truncate(2)}s"
+          else
+            "#{(self * 1000).to_i}ms"
+          end
+        end
+      end
+    end
+  end
+end

data/lib/.rbnext/3.0/wsdirector/protocols/base.rb

@@ -0,0 +1,288 @@
+# frozen_string_literal: true
+
+require "time"
+require "wsdirector/ext/formatting"
+
+module WSDirector
+  module Protocols
+    using Ext::Formatting
+
+    using(Module.new do
+      refine ::Object do
+        def matches?(other)
+          self == other
+        end
+
+        def partially_matches?(other)
+          self == other
+        end
+      end
+
+      refine ::Array do
+        def matches?(actual)
+          if actual.is_a?(::String)
+            actual = JSON.parse(actual) rescue nil # rubocop:disable Style/RescueModifier
+          end
+
+          return false unless actual
+
+          each.with_index do
+            return false unless _1.matches?(actual[_2])
+          end
+
+          true
+        end
+
+        def partially_matches?(actual)
+          if actual.is_a?(::String)
+            actual = JSON.parse(actual) rescue nil # rubocop:disable Style/RescueModifier
+          end
+
+          return false unless actual
+
+          each.with_index do
+            return false unless _1.partially_matches?(actual[_2])
+          end
+
+          true
+        end
+      end
+
+      refine ::Hash do
+        def matches?(actual)
+          unless actual.is_a?(::Hash)
+            actual = JSON.parse(actual) rescue nil # rubocop:disable Style/RescueModifier
+          end
+
+          return false unless actual
+
+          actual.each_key do
+            return false unless actual[_1].matches?(self[_1])
+          end
+
+          true
+        end
+
+        def partially_matches?(actual)
+          unless actual.is_a?(::Hash)
+            actual = JSON.parse(actual) rescue nil # rubocop:disable Style/RescueModifier
+          end
+
+          return false unless actual
+
+          each_key do
+            return false unless self[_1].partially_matches?(actual[_1])
+          end
+
+          true
+        end
+      end
+    end)
+
+    class PartialMatcher
+      attr_reader :obj
+
+      def initialize(obj)
+        @obj = obj
+      end
+
+      def matches?(actual)
+        obj.partially_matches?(actual)
+      end
+
+      def inspect
+        "an object including #{obj.inspect}"
+      end
+
+      def truncate(...) ;  obj.truncate(...); end
+    end
+
+    # Base protocol describes basic actions
+    class Base
+      include WSDirector::Utils
+
+      def initialize(task, scale: 1, logger: nil, id: nil, color: nil)
+        @task = task
+        @scale = scale
+        @logger = logger
+        @id = id
+        @color = color
+      end
+
+      def init_client(...)
+        log { "Connecting" }
+
+        @client = build_client(...)
+
+        log(:done) { "Connected" }
+      end
+
+      def handle_step(step)
+        type = step.delete("type")
+        raise Error, "Unknown step: #{type}" unless respond_to?(type)
+
+        return unless task.sampled?(step)
+
+        public_send(type, step)
+      end
+
+      # Sleeps for a specified number of seconds.
+      #
+      # If "shift" is provided than the initial value is
+      # shifted by random number from (-shift, shift).
+      #
+      # Set "debug" to true to print the delay time.
+      def sleep(step)
+        delay = step.fetch("time").to_f
+        shift = step.fetch("shift", 0).to_f
+
+        delay = delay - shift * rand + shift * rand
+
+        log { "Sleep for #{delay}s" }
+
+        Kernel.sleep delay if delay > 0
+
+        log(:done) { "Slept for #{delay}s" }
+      end
+
+      # Prints provided message
+      def debug(step)
+        with_logger do
+          log(nil) { step.fetch("message") }
+        end
+      end
+
+      def receive(step)
+        expected = step["data"] || PartialMatcher.new(step["data>"])
+        ordered = step["ordered"]
+
+        log { "Receive a message: #{expected.truncate(50)}" }
+
+        received = nil
+
+        client.each_message do |msg, id|
+          received = msg
+          if expected.matches?(msg)
+            client.consumed(id)
+            break
+          end
+
+          if ordered
+            raise UnmatchedExpectationError, prepare_receive_error(expected, received)
+          end
+        end
+
+        log(:done) { "Received a message: #{received&.truncate(50)}" }
+      rescue ThreadError
+        if received
+          raise UnmatchedExpectationError, prepare_receive_error(expected, received)
+        else
+          raise NoMessageError, "Expected to receive #{expected} but nothing has been received"
+        end
+      end
+
+      def receive_all(step)
+        messages = step.delete("messages")
+        raise ArgumentError, "Messages array must be specified" if
+          messages.nil? || messages.empty?
+
+        expected =
+          messages.map do |msg|
+            multiplier = parse_multiplier(msg.delete("multiplier") || "1")
+            [msg["data"] || PartialMatcher.new(msg["data>"]), multiplier]
+          end.to_h
+
+        total_expected = expected.values.sum
+        total_received = 0
+
+        log { "Receive #{total_expected} messages" }
+
+        total_expected.times do
+          received = client.receive
+
+          total_received += 1
+
+          match = expected.find { |k, _| k.matches?(received) }
+
+          raise UnexpectedMessageError, "Unexpected message received: #{received}" if
+            match.nil?
+
+          expected[match.first] -= 1
+          expected.delete(match.first) if expected[match.first].zero?
+        end
+
+        log(:done) { "Received #{total_expected} messages" }
+      rescue ThreadError
+        raise NoMessageError,
+          "Expected to receive #{total_expected} messages " \
+          "but received only #{total_received}"
+      end
+      # rubocop: enable Metrics/CyclomaticComplexity
+
+      def send(step)
+        data = step.fetch("data")
+        data = JSON.generate(data) if data.is_a?(Hash)
+
+        client.send(data)
+
+        log(nil) { "Sent message: #{data.truncate(50)}" }
+      end
+
+      def wait_all(_step)
+        log { "Wait all clients" }
+        task.global_holder.wait_all
+        log { "All clients" }
+      end
+
+      def to_proc
+        proc { |step| handle_step(step) }
+      end
+
+      private
+
+      attr_reader :client, :task, :logger, :id, :color
+
+      def build_client(...)
+        Client.new(...)
+      end
+
+      def prepare_receive_error(expected, received)
+        <<~MSG
+          Action failed: #receive
+             -- expected: #{expected.inspect}
+             ++ got: #{received} (#{received.class})
+        MSG
+      end
+
+      def log(state = :begin)
+        return unless logger
+
+        if state == :begin
+          @last_event_at = Time.now.to_f
+        end
+
+        done_info =
+          if state == :done
+            " (#{(Time.now.to_f - @last_event_at).duration})"
+          else
+            ""
+          end
+
+        msg = "[#{Time.now.strftime("%H:%I:%S.%L")}] client=#{id} #{yield}#{done_info}"
+
+        msg = msg.colorize(color) if color
+
+        logger.puts msg
+      end
+
+      def with_logger
+        return yield if logger
+
+        @logger = $stdout
+        yield
+      ensure
+        @logger = nil
+      end
+    end
+  end
+end

data/lib/.rbnext/3.1/wsdirector/cli.rb

@@ -0,0 +1,152 @@
+# frozen_string_literal: true
+
+require "optparse"
+require "uri"
+
+require "wsdirector"
+
+module WSDirector
+  # Command line interface for WsDirector
+  class CLI
+    class Configuration
+      attr_accessor :ws_url, :scenario_path, :colorize, :scale,
+        :sync_timeout, :json_scenario, :subprotocol, :verbose
+
+      def initialize
+        reset!
+      end
+
+      def colorize?
+        colorize == true
+      end
+
+      # Restore to defaults
+      def reset!
+        @scale = 1
+        @colorize = $stdout.tty?
+        @sync_timeout = 5
+      end
+    end
+
+    attr_reader :config
+
+    def run
+      @config = Configuration.new
+
+      parse_args!
+
+      begin
+        require "colorize" if config.colorize?
+      rescue LoadError
+        config.colorize = false
+        warn "Install colorize to use colored output"
+      end
+
+      connection_options = {
+        subprotocol: config.subprotocol
+      }.compact
+
+      scenario = config.scenario_path || config.json_scenario
+
+      config.ws_url = "ws://#{config.ws_url}" unless config.ws_url.start_with?(/wss?:\/\//)
+
+      url = config.ws_url
+      scale = config.scale
+      sync_timeout = config.sync_timeout
+      colorize = config.colorize
+
+      logger = $stdout if config.verbose
+
+      result = WSDirector.run(
+        scenario,
+        url: url,
+        connection_options: connection_options,
+        scale: scale,
+        sync_timeout: sync_timeout,
+        logger: logger,
+        colorize: colorize
+      )
+
+      puts "\n\n" if config.verbose
+
+      result.print_summary(colorize: config.colorize?)
+      result.success? || exit(1)
+    end
+
+    private
+
+    FILE_FORMAT = /.+.(json|yml)\z/
+    private_constant :FILE_FORMAT
+
+    def parse_args!
+      parser = OptionParser.new do |opts|
+        opts.banner = "Usage: wsdirector scenario_path ws_url [options]"
+
+        opts.on("-s SCALE", "--scale=SCALE", Integer, "Scale factor") do
+          config.scale = _1
+        end
+
+        opts.on("-t TIMEOUT", "--timeout=TIMEOUT", Integer, "Synchronization (wait_all) timeout") do
+          config.sync_timeout = _1
+        end
+
+        opts.on("-i JSON", "--include=JSON", String, "Include JSON to parse") do
+          config.json_scenario = _1
+        end
+
+        opts.on("-u URL", "--url=URL", Object, "Websocket server URL") do
+          config.ws_url = _1
+        end
+
+        opts.on("-f PATH", "--file=PATH", String, "Scenario path") do
+          config.scenario_path = _1
+        end
+
+        opts.on("--subprotocol=VALUE", String, "WebSocket subprotocol") do
+          config.subprotocol = _1
+        end
+
+        opts.on("-c", "--[no-]color", "Colorize output") do
+          config.colorize = _1
+        end
+
+        opts.on("-v", "--version", "Print version") do
+          $stdout.puts WSDirector::VERSION
+          exit 0
+        end
+
+        opts.on("-vv", "Print verbose logs") do
+          config.verbose = true
+        end
+
+        opts.on("-r", "--require=PATH", "Load Ruby file (e.g., protocol)") do
+          Kernel.load(_1)
+        end
+      end
+
+      parser.parse!
+
+      unless config.scenario_path
+        config.scenario_path = ARGV.grep(FILE_FORMAT).last
+      end
+
+      unless config.ws_url
+        config.ws_url = ARGV.grep(URI::DEFAULT_PARSER.make_regexp).last
+      end
+
+      check_for_errors
+    end
+
+    def check_for_errors
+      if config.json_scenario.nil?
+        raise(Error, "Scenario is missing") unless config.scenario_path
+
+        unless File.file?(config.scenario_path)
+          raise(Error, "File doesn't exist # config.scenario_path}")
+        end
+      end
+
+      raise(Error, "Websocket server url is missing") unless config.ws_url
+    end
+  end
+end