#!/usr/bin/env ruby

#---
# Copyright 2003, 2004, 2005 by Jim Weirich (jim@weirichhouse.org).
# All rights reserved.

# Permission is granted for use, copying, modification, distribution,
# and distribution of modified versions of this work as long as the
# above copyright notice is included.
#+++

require 'test/unit'

######################################################################
# FlexMock is a flexible mock object suitable for using with Ruby's
# Test::Unit unit test framework.  FlexMock has a simple interface
# that's easy to remember, and leaves the hard stuff to all those
# other mock object implementations.
#
# Basic Usage:
#
#   m = FlexMock.new("name")
#   m.mock_handle(:meth) { |args| assert_stuff }
#
# Simplified Usage:
#
#   m = FlexMock.new("name")
#   m.should_receive(:upcase).with("stuff").
#     returns("STUFF")
#   m.should_receive(:downcase).with(String).
#     returns { |s| s.downcase }.once
#
class FlexMock
  include Test::Unit::Assertions

  attr_reader :mock_name, :mock_groups
  attr_accessor :mock_current_order

  # Create a FlexMock object with the given name.  The name is used in
  # error messages.
  def initialize(name="unknown")
    @mock_name = name
    @expectations = Hash.new
    @allocated_order = 0
    @mock_current_order = 0
    @mock_groups = {}
  end
  
  # Handle all messages denoted by +sym+ by calling the given block
  # and passing any parameters to the block.  If we know exactly how
  # many calls are to be made to a particular method, we may check
  # that by passing in the number of expected calls as a second
  # paramter.
  def mock_handle(sym, expected_count=nil, &block)
    self.should_receive(sym).times(expected_count).returns(&block)
  end

  # Verify that each method that had an explicit expected count was
  # actually called that many times.
  def mock_verify
    mock_wrap do
      @expectations.each do |sym, handler|
	handler.mock_verify
      end
    end
  end

  # Allocation a new order number from the mock.
  def mock_allocate_order
    @auto_allocate = true
    @allocated_order += 1
  end

  # Ignore all undefined (missing) method calls.
  def should_ignore_missing
    @ignore_missing = true
  end
  alias mock_ignore_missing should_ignore_missing

  # Handle missing methods by attempting to look up a handler.
  def method_missing(sym, *args, &block)
    mock_wrap do
      if handler = @expectations[sym]
	args << block  if block_given?
	handler.call(*args)
      else
	super(sym, *args, &block)  unless @ignore_missing
      end
    end
  end

  # Override the built-in respond_to? to include the mocked methods.
  def respond_to?(sym)
    super || (@expectations[sym] ? true : @ignore_missing)
  end

  # Override the built-in +method+ to include the mocked methods.
  def method(sym)
    @expectations[sym] || super
  rescue NameError => ex
    if @ignore_missing
      proc { }
    else
      raise ex
    end
  end

  # Declare that the mock object should receive a message with the
  # given name.  An expectation object for the method name is returned
  # as the result of this method.  Further expectation constraints can
  # be added by chaining to the result.
  #
  # See Expectation for a list of declarators that can be used.
  def should_receive(sym)
    @expectations[sym] ||= ExpectationDirector.new(sym)
    result = Expectation.new(self, sym)
    @expectations[sym] << result
    result
  end

  class << self
    include Test::Unit::Assertions

    # Class method to make sure that verify is called at the end of a
    # test.  One mock object will be created for each name given to
    # the use method.  The mocks will be passed to the block as
    # arguments.  If no names are given, then a single anonymous mock
    # object will be created.
    #
    # At the end of the use block, each mock object will be verified
    # to make sure the proper number of calls have been made.
    #
    # Usage:
    #
    #   FlexMock.use("name") do |mock|    # Creates a mock named "name"
    #     mock.should_receive(:meth).
    #       returns(0).once
    #   end                               # mock is verified here
    #
    def use(*names)
      names = ["unknown"] if names.empty?
      got_excecption = false
      mocks = names.collect { |n| new(n) }
      yield(*mocks)
    rescue Exception => ex
      got_exception = true
      raise
    ensure
      mocks.each do |mock|
	mock.mock_verify     unless got_exception
      end
    end
    
    # Class method to format a method name and argument list as a nice
    # looking string.
    def format_args(sym, args)
      if args
	"#{sym}(#{args.collect { |a| a.inspect }.join(', ')})"
      else
	"#{sym}(*args)"
      end
    end
    
    # Check will assert the block returns true.  If it doesn't, an
    # assertion failure is triggered with the given message.
    def check(msg, &block)
      assert_block(msg, &block)
    end
  end

  private

  # Wrap a block of code so the any assertion errors are wrapped so
  # that the mock name is added to the error message .
  def mock_wrap(&block)
    yield
  rescue Test::Unit::AssertionFailedError => ex
    raise Test::Unit::AssertionFailedError, 
      "in mock '#{@mock_name}': #{ex.message}",
      ex.backtrace
  end

  ####################################################################
  # The expectation director is responsible for routing calls to the
  # correct expectations for a given argument list.
  #
  class ExpectationDirector

    # Create an ExpectationDirector for a mock object.
    def initialize(sym)
      @sym = sym
      @expectations = []
      @expected_order = nil
    end

    # Invoke the expectations for a given set of arguments.
    def call(*args)
      exp = @expectations.find { |e| e.match_args(args) } ||
	@expectations.find { |e| e.expected_args.nil? }
      FlexMock.check("no matching handler found for " +
 	FlexMock.format_args(@sym, args)) { ! exp.nil? }
      exp.verify_call(*args)
    end

    # Same as call.
    def [](*args)
      call(*args)
    end

    # Append an expectation to this director.
    def <<(expectation)
      @expectations << expectation
    end

    # Do the post test verification for this directory.  Check all the
    # expectations.
    def mock_verify
      @expectations.each do |exp|
	exp.mock_verify
      end
    end
  end

  ####################################################################
  # Match any object
  class AnyMatcher
    def ===(target)
      true
    end
    def inspect
      "ANY"
    end
  end
  
  ####################################################################
  # Match only things that are equal.
  class EqualMatcher
    def initialize(obj)
      @obj = obj
    end
    def ===(target)
      @obj == target
    end
    def inspect
      "==(#{@obj.inspect})"
    end
  end
  
  ANY = AnyMatcher.new

  ####################################################################
  # Include this module in your test class if you wish to use the +eq+
  # and +any+ argument matching methods without a prefix.  (Otherwise
  # use <tt>FlexMock.any</tt> and <tt>FlexMock.eq(obj)</tt>.
  #
  module ArgumentTypes
    # Return an argument matcher that matches any argument.
    def any
      ANY
    end

    # Return an argument matcher that only matches things equal to
    # (==) the given object.
    def eq(obj)
      EqualMatcher.new(obj)
    end
  end

  ####################################################################
  # Base class for all the count validators.
  #
  class CountValidator
    include Test::Unit::Assertions
    def initialize(expectation, limit)
      @exp = expectation
      @limit = limit
    end
  end

  ####################################################################
  # Validator for exact call counts.
  #
  class ExactCountValidator < CountValidator
    # Validate that the method expectation was called exactly +n+
    # times.
    def validate(n)
      assert_equal @limit, n, 
	"method '#{@exp}' called incorrect number of times"
    end
  end

  ####################################################################
  # Validator for call counts greater than or equal to a limit.
  #
  class AtLeastCountValidator < CountValidator
    # Validate the method expectation was called no more than +n+
    # times.
    def validate(n)
      assert n >= @limit,
	"Method '#{@exp}' should be called at least #{@limit} times,\n" +
	"only called #{n} times"
    end
  end

  ####################################################################
  # Validator for call counts less than or equal to a limit.
  #
  class AtMostCountValidator < CountValidator
    # Validate the method expectation was called at least +n+ times.
    def validate(n)
      assert n <= @limit,
	"Method '#{@exp}' should be called at most #{@limit} times,\n" +
	"only called #{n} times"
    end
  end

  ####################################################################
  # An Expectation is returned from each +should_receive+ message sent
  # to mock object.  Each expectation records how a message matching
  # the message name (argument to +should_receive+) and the argument
  # list (given by +with+) should behave.  Mock expectations can be
  # recorded by chaining the declaration methods defined in this
  # class.
  #
  # For example:
  #
  #   mock.should_receive(:meth).with(args).and_returns(result)
  #
  class Expectation
    include Test::Unit::Assertions

    attr_reader :expected_args, :mock, :order_number

    # Create an expectation for a method named +sym+.
    def initialize(mock, sym)
      @mock = mock
      @sym = sym
      @expected_args = nil
      @count_validators = []
      @count_validator_class = ExactCountValidator
      @actual_count = 0
      @return_value = nil
      @return_block = lambda { @return_value }
      @order_number = nil
    end

    def to_s
      FlexMock.format_args(@sym, @expected_args)
    end

    # Verify the current call with the given arguments matches the
    # expectations recorded in this object.
    def verify_call(*args)
      validate_order
      @actual_count += 1
      @return_block.call(*args)
    end

    # Validate that the order 
    def validate_order
      return if @order_number.nil?
      FlexMock.check("method #{to_s} called out of order " +
	"(expected order #{@order_number}, was #{@mock.mock_current_order})") {
	@order_number >= @mock.mock_current_order
      }
      @mock.mock_current_order = @order_number
    end
    private :validate_order

    # Validate the correct number of calls have been made.  Called by
    # the teardown process.
    def mock_verify
      @count_validators.each do |v|
	v.validate(@actual_count)
      end
    end

    # Does the argument list match this expectation's argument
    # specification.
    def match_args(args)
      return false if @expected_args.nil?
      return false if args.size != @expected_args.size
      (0...args.size).all? { |i| match_arg(@expected_args[i], args[i]) }
    end

    # Does the expected argument match the corresponding actual value.
    def match_arg(expected, actual)
      expected === actual ||
	expected == actual ||
	( Regexp === expected && expected === actual.to_s )
    end

    # Declare that the method should expect the given argument list.
    def with(*args)
      @expected_args = args
      self
    end

    # Declare that the method should be called with no arguments.
    def with_no_args
      with
    end

    # Declare that the method can be called with any number of
    # arguments of any type.
    def with_any_args
      @expected_args = nil
      self
    end

    # Declare that the method returns a particular value (when the
    # argument list is matched).  If a block is given, it is evaluated
    # on each call and its value is returned.  +and_return+ is an
    # alias for +returns+.
    #
    # For example:
    # 
    #  mock.should_receive(:f).returns(12)   # returns 12
    #
    #  mock.should_receive(:f).with(String). # returns an
    #    returns { |str| str.upcase }        # upcased string
    #
    def returns(value=nil, &block)
      @return_block = block_given? ? block : lambda { value }
      self
    end
    alias :and_return :returns	# :nodoc:

    # Declare that the method may be called any number of times.
    def zero_or_more_times
      at_least.never
    end

    # Declare that the method is called +limit+ times with the
    # declared argument list.  This may be modified by the +at_least+
    # and +at_most+ declarators.
    def times(limit)
      @count_validators << @count_validator_class.new(self, limit) unless limit.nil?
      @count_validator_class = ExactCountValidator
      self
    end

    # Declare that the method is never expected to be called with the
    # given argument list.  This may be modified by the +at_least+ and
    # +at_most+ declarators.
    def never
      times(0)
    end

    # Declare that the method is expected to be called exactly once
    # with the given argument list.  This may be modified by the
    # +at_least+ and +at_most+ declarators.
    def once
      times(1)
    end

    # Declare that the method is expected to be called exactly twice
    # with the given argument list.  This may be modified by the
    # +at_least+ and +at_most+ declarators.
    def twice
      times(2)
    end

    # Modifies the next call count declarator (+times+, +never+,
    # +once+ or +twice+) so that the declarator means the method is
    # called at least that many times.
    #
    # E.g. method f must be called at least twice:
    #
    #   mock.should_receive(:f).at_least.twice
    #
    def at_least
      @count_validator_class = AtLeastCountValidator
      self
    end

    # Modifies the next call count declarator (+times+, +never+,
    # +once+ or +twice+) so that the declarator means the method is
    # called at most that many times.
    #
    # E.g. method f must be called no more than twice
    #
    #   mock.should_receive(:f).at_most.twice
    #
    def at_most
      @count_validator_class = AtMostCountValidator
      self
    end

    # Declare that the given method must be called in order.  All
    # ordered method calls must be received in the order specified by
    # the ordering of the +should_receive+ messages.  Receiving a
    # methods out of the specified order will cause a test failure.
    #
    # If the user needs more fine control over ordering
    # (e.g. specifying that a group of messages may be received in any
    # order as long as they all come after another group of messages),
    # a _group_ _name_ may be specified in the +ordered+ calls.  All
    # messages within the same group may be recieved in any order.
    #
    # For example, in the following, messages +flip+ and +flop+ may be
    # received in any order (because they are in the same group), but
    # must occur strictly after +start+ but before +end+.  The message
    # +any_time+ may be received at any time because it is not
    # ordered.
    #
    #    m = FlexMock.new
    #    m.should_receive(:any_time)
    #    m.should_receive(:start).ordered
    #    m.should_receive(:flip).ordered(:flip_flop_group)
    #    m.should_receive(:flop).ordered(:flip_flop_group)
    #    m.should_receive(:end).ordered
    #
    def ordered(group_name=nil)
      if group_name.nil?
	@order_number = @mock.mock_allocate_order
      elsif (num = @mock.mock_groups[group_name])
	@order_number = num
      else
	@order_number = @mock.mock_allocate_order
	@mock.mock_groups[group_name] = @order_number
      end
      self
    end
  end

  extend ArgumentTypes
end