Merge pull request #8 from ohbarye/choose-one_of-arbitrary

Add ChooseArbitrary and OneOfArbitrary

lib/pbt/arbitrary/arbitrary_methods.rb

@@ -4,6 +4,8 @@
 require "pbt/arbitrary/integer_arbitrary"
 require "pbt/arbitrary/tuple_arbitrary"
 require "pbt/arbitrary/fixed_hash_arbitrary"
+require "pbt/arbitrary/choose_arbitrary"
+require "pbt/arbitrary/one_of_arbitrary"
 
 module Pbt
   module Arbitrary
@@ -38,6 +40,16 @@ def tuple(*arbs)
       def fixed_hash(hash)
         FixedHashArbitrary.new(hash)
       end
+
+      # @param range [Range<Integer>]
+      def choose(range)
+        ChooseArbitrary.new(range)
+      end
+
+      # @param choices [Array]
+      def one_of(*choices)
+        OneOfArbitrary.new(choices)
+      end
     end
   end
 end

lib/pbt/arbitrary/choose_arbitrary.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Pbt
+  module Arbitrary
+    class ChooseArbitrary
+      # @param range [Range<Integer>]
+      def initialize(range)
+        @range = range
+      end
+
+      # @return [Array]
+      def generate(rng)
+        rng.rand(@range)
+      end
+
+      # range is ordered from min to max, so we can just shrink towards min
+      # @return [Enumerator]
+      def shrink(current)
+        min, max = [@range.begin, @range.end].sort
+        IntegerArbitrary.new(min, max).shrink(current, target: min)
+      end
+    end
+  end
+end

lib/pbt/arbitrary/one_of_arbitrary.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Pbt
+  module Arbitrary
+    class OneOfArbitrary
+      # @param choices [Array]
+      def initialize(choices)
+        @choices = choices
+        @idx_arb = IntegerArbitrary.new(0, choices.size - 1)
+      end
+
+      # @return [Array]
+      def generate(rng)
+        @choices[@idx_arb.generate(rng)]
+      end
+
+      # Shrinks to values earlier in the list of `choices`.
+      # @return [Enumerator]
+      def shrink(current)
+        Enumerator.new do |y|
+          @idx_arb.shrink(@choices.index(current)).map do |idx|
+            y << @choices[idx]
+          end
+        end
+      end
+    end
+  end
+end

spec/pbt/arbitrary/choose_arbitrary_spec.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+RSpec.describe Pbt::Arbitrary::ChooseArbitrary do
+  describe "#generate" do
+    it "generates an integer" do
+      val = Pbt::Arbitrary::ChooseArbitrary.new(1..10).generate(Random.new)
+      expect(val).to be_a(Integer)
+    end
+
+    it "generates an integer in given range" do
+      aggregate_failures do
+        100.times do
+          val = Pbt::Arbitrary::ChooseArbitrary.new(-1..4).generate(Random.new)
+          expect(val).to be >= -1
+          expect(val).to be <= 4
+        end
+      end
+    end
+  end
+
+  describe "#shrink" do
+    it "returns an Enumerator" do
+      arb = Pbt::Arbitrary::ChooseArbitrary.new(1..10)
+      val = arb.generate(Random.new)
+      expect(arb.shrink(val)).to be_a(Enumerator)
+    end
+
+    it "returns an Enumerator that iterates halved integers towards the min" do
+      arb = Pbt::Arbitrary::ChooseArbitrary.new(-2..10)
+      expect(arb.shrink(50).to_a).to eq [24, 11, 5, 2, 0, -1, -2]
+    end
+
+    context "when current value and target is same" do
+      it "returns an empty Enumerator" do
+        arb = Pbt::Arbitrary::ChooseArbitrary.new(-2..10)
+        expect(arb.shrink(-2).to_a).to eq []
+      end
+    end
+  end
+end

spec/pbt/arbitrary/one_of_arbitrary_spec.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+RSpec.describe Pbt::Arbitrary::OneOfArbitrary do
+  describe "#generate" do
+    it "generates any integer in given array" do
+      aggregate_failures do
+        choices = [:a, 1, "A", 2.0, {foo: :bar}, [1], String]
+        100.times do
+          val = Pbt::Arbitrary::OneOfArbitrary.new(choices).generate(Random.new)
+          expect(choices).to include val
+        end
+      end
+    end
+  end
+
+  describe "#shrink" do
+    it "returns an Enumerator" do
+      arb = Pbt::Arbitrary::OneOfArbitrary.new([1])
+      val = arb.generate(Random.new)
+      expect(arb.shrink(val)).to be_a(Enumerator)
+    end
+
+    it "returns an Enumerator that iterates halved integers towards the min" do
+      arb = Pbt::Arbitrary::OneOfArbitrary.new([:z, :a, :x, :y, :c, :d])
+      expect(arb.shrink(:c).to_a).to eq [:x, :a, :z]
+    end
+
+    context "when current value and target is same" do
+      it "returns an empty Enumerator" do
+        arb = Pbt::Arbitrary::OneOfArbitrary.new(["A"])
+        expect(arb.shrink("A").to_a).to eq []
+      end
+    end
+  end
+end