kazoo-ruby 0.3.3 → 0.4.0

data/lib/kazoo/version.rb

@@ -1,3 +1,3 @@
 module Kazoo
-  VERSION = "0.3.3"
+  VERSION = "0.4.0"
 end

data/test/functional/functional_consumergroup_test.rb

@@ -7,11 +7,18 @@ class FunctionalConsumergroupTest < Minitest::Test
     @cg = Kazoo::Consumergroup.new(@cluster, 'test.kazoo')
     @cg.create
 
-    @topic = @cluster.create_topic("test.4", partitions: 4, replication_factor: 1)
+    @topic1 = @cluster.topic('test.1')
+    @topic1.destroy if @topic1.exists?
+    @topic1 = @cluster.create_topic('test.1', partitions: 1, replication_factor: 1)
+
+    @topic4 = @cluster.topic('test.4')
+    @topic4.destroy if @topic4.exists?
+    @topic4 = @cluster.create_topic('test.4', partitions: 4, replication_factor: 1)
   end
 
   def teardown
-    @topic.destroy
+    @topic1.destroy if @topic1.exists?
+    @topic4.destroy if @topic4.exists?
 
     cg = Kazoo::Consumergroup.new(@cluster, 'test.kazoo')
     cg.destroy if cg.exists?
@@ -30,32 +37,19 @@ class FunctionalConsumergroupTest < Minitest::Test
     refute cg.exists?
   end
 
-  def test_retrieve_and_commit_offsets
-    topic = Kazoo::Topic.new(@cluster, 'test.1')
-    partition = Kazoo::Partition.new(topic, 0)
-
-    assert_nil @cg.retrieve_offset(partition)
-
-    @cg.commit_offset(partition, 1234)
-
-    assert_equal 1234 + 1, @cg.retrieve_offset(partition)
-
-    @cg.reset_all_offsets
-    assert_nil @cg.retrieve_offset(partition)
-  end
-
   def test_unclaimed_partitions
-    partition40 = @topic.partition(0)
-    partition41 = @topic.partition(1)
-    partition42 = @topic.partition(2)
-    partition43 = @topic.partition(3)
+    partition40 = @topic4.partition(0)
+    partition41 = @topic4.partition(1)
+    partition42 = @topic4.partition(2)
+    partition43 = @topic4.partition(3)
 
     refute @cg.active?
 
-    instance1 = @cg.instantiate
-    instance1.register([@topic])
-    instance2 = @cg.instantiate
-    instance2.register([@topic])
+    subscription = Kazoo::Subscription.build(@topic4)
+    instance1 = @cg.instantiate(subscription: subscription)
+    instance1.register
+    instance2 = @cg.instantiate(subscription: subscription)
+    instance2.register
 
     assert @cg.active?
 
@@ -63,60 +57,118 @@ class FunctionalConsumergroupTest < Minitest::Test
     instance2.claim_partition(partition41)
 
     assert_equal 2, @cg.partition_claims.length
-    assert_equal [@topic], @cg.topics
-    assert_equal @topic.partitions, @cg.partitions
+    assert_equal [@topic4], @cg.topics
+    assert_equal @topic4.partitions, @cg.partitions
 
     assert_equal Set[partition42, partition43], Set.new(@cg.unclaimed_partitions)
   end
 
-  def test_retrieve_all_offsets
-    topic1 = Kazoo::Topic.new(@cluster, 'test.1')
-    partition10 = Kazoo::Partition.new(topic1, 0)
+  def test_watch_instances
+    instance1 = @cg.instantiate(subscription: @topic1).register
+    instance2 = @cg.instantiate(subscription: @topic1).register
 
-    topic4 = Kazoo::Topic.new(@cluster, 'test.4')
-    partition40 = Kazoo::Partition.new(topic4, 0)
-    partition41 = Kazoo::Partition.new(topic4, 1)
-    partition42 = Kazoo::Partition.new(topic4, 2)
-    partition43 = Kazoo::Partition.new(topic4, 3)
+    t = Thread.current
+    instances, cb = @cg.watch_instances { t.run if t.status == 'sleep' }
+    assert_equal Set[instance1, instance2], Set.new(instances)
 
-    assert_equal Hash.new, @cg.retrieve_all_offsets
+    Thread.new { instance2.deregister }
+
+    Thread.stop unless cb.completed?
+
+    assert assert_equal Set[instance1], Set.new(@cg.instances)
+    instance1.deregister
+  end
+
+  def test_clean_topic_claims
+    deleted_topic = @cluster.topic('non_existing')
+    deleted_topic.destroy if deleted_topic.exists?
+    deleted_topic = @cluster.create_topic('non_existing', partitions: 1, replication_factor: 1)
+
+    subscription = Kazoo::Subscription.build([@topic4, deleted_topic])
+    @cg.instantiate(subscription: subscription).register
+
+    deleted_topic.destroy
+
+    assert_equal Set[@topic4, deleted_topic], Set.new(@cg.claimed_topics)
+
+    @cg.clean_topic_claims(Kazoo::Subscription.everything)
+    assert_equal Set[@topic4], Set.new(@cg.claimed_topics)
+
+    @cg.clean_topic_claims(Kazoo::Subscription.build([]))
+    assert_equal [], @cg.claimed_topics
+  end
+
+  def test_retrieve_and_commit_offsets
+    partition = Kazoo::Partition.new(@topic1, 0)
+
+    assert_nil @cg.retrieve_offset(partition)
+
+    @cg.commit_offset(partition, 1234)
+
+    assert_equal 1234 + 1, @cg.retrieve_offset(partition)
+
+    @cg.reset_all_offsets
+    assert_nil @cg.retrieve_offset(partition)
+  end
+
+  def test_retrieve_offsets_and_reset_all_offsets
+    partition10 = Kazoo::Partition.new(@topic1, 0)
+    partition40 = Kazoo::Partition.new(@topic4, 0)
+    partition41 = Kazoo::Partition.new(@topic4, 1)
+    partition42 = Kazoo::Partition.new(@topic4, 2)
+    partition43 = Kazoo::Partition.new(@topic4, 3)
+
+    assert @cg.retrieve_offsets(Kazoo::Subscription.everything).values.all? { |v| v.nil? }
 
     @cg.commit_offset(partition10, 10)
     @cg.commit_offset(partition40, 40)
     @cg.commit_offset(partition41, 41)
     @cg.commit_offset(partition42, 42)
-    @cg.commit_offset(partition43, 43)
 
-    offsets = @cg.retrieve_all_offsets
+    offsets = @cg.retrieve_offsets(Kazoo::Subscription.everything)
 
-    assert_equal 5, offsets.length
+    assert offsets.length >= 5
     assert_equal 11, offsets[partition10]
     assert_equal 41, offsets[partition40]
     assert_equal 42, offsets[partition41]
     assert_equal 43, offsets[partition42]
-    assert_equal 44, offsets[partition43]
+    assert_equal nil, offsets[partition43]
+
+    offsets = @cg.retrieve_offsets(['test.1'])
+
+    assert_equal 1, offsets.length
+    assert_equal 11, offsets[partition10]
 
     @cg.reset_all_offsets
-    assert_equal Hash.new, @cg.retrieve_all_offsets
+    assert @cg.retrieve_offsets(Kazoo::Subscription.everything).values.all? { |v| v.nil? }
   end
 
-  def test_watch_instances
-    topic = Kazoo::Topic.new(@cluster, 'test.1')
+  def test_retrieve_all_offsets_and_cleanup_offsets
+    deleted_topic = @cluster.topic('non_existing')
+    deleted_topic.destroy if deleted_topic.exists?
+    deleted_topic = @cluster.create_topic('non_existing', partitions: 1, replication_factor: 1)
 
-    instance1 = @cg.instantiate
-    instance1.register([topic])
-    instance2 = @cg.instantiate
-    instance2.register([topic])
+    subscription = Kazoo::Subscription.build([@topic4, deleted_topic])
 
-    t = Thread.current
-    instances, cb = @cg.watch_instances { t.run if t.status == 'sleep' }
-    assert_equal Set[instance1, instance2], Set.new(instances)
+    @cg.commit_offset(@topic4.partition(0), 1234)
+    @cg.commit_offset(deleted_topic.partition(0), 1234)
 
-    Thread.new { instance2.deregister }
+    deleted_topic.destroy
 
-    Thread.stop unless cb.completed?
+    expected_offsets = {
+      deleted_topic.partition(0) => 1234 + 1,
+      @topic4.partition(0) => 1234 + 1,
+    }
+    assert_equal expected_offsets, @cg.retrieve_all_offsets
 
-    assert assert_equal Set[instance1], Set.new(@cg.instances)
-    instance1.deregister
+    @cg.clean_stored_offsets(Kazoo::Subscription.everything)
+
+    expected_offsets = {
+      @topic4.partition(0) => 1234 + 1,
+    }
+    assert_equal expected_offsets, @cg.retrieve_all_offsets
+
+    @cg.clean_stored_offsets([])
+    assert_equal Hash.new, @cg.retrieve_all_offsets
   end
 end

data/test/functional/functional_subscription_test.rb

@@ -0,0 +1,62 @@
+require 'test_helper'
+
+class FunctionalTopicManagementTest < Minitest::Test
+  def setup
+    zookeeper = ENV["ZOOKEEPER_PEERS"] || "127.0.0.1:2181"
+    @cluster = Kazoo.connect(zookeeper)
+  end
+
+  def test_consumergroup_subscription
+    cg = @cluster.consumergroup('test.consumergroup.subscription')
+    cg.create
+
+    assert_raises(Kazoo::NoRunningInstances) { cg.subscription }
+
+    subscription_1  = Kazoo::Subscription.build('topic.1')
+    subscription_14 = Kazoo::Subscription.build(['topic.1', 'topic.4'])
+
+    instance1 = cg.instantiate(subscription: subscription_1).register
+    instance2 = cg.instantiate(subscription: subscription_1).register
+    instance3 = cg.instantiate(subscription: subscription_14).register
+
+    assert_raises(Kazoo::InconsistentSubscriptions) { cg.subscription }
+
+    instance3.deregister
+
+    assert_equal subscription_1, cg.subscription
+    assert_equal subscription_1, instance1.subscription
+
+    instance2.deregister
+
+    assert_equal subscription_1, cg.subscription
+    assert_equal subscription_1, instance1.subscription
+
+    instance1.deregister
+
+    assert_raises(Kazoo::NoRunningInstances) { cg.subscription }
+  ensure
+    cg.destroy if cg.exists?
+  end
+
+  def test_subscription_topics_and_partitions
+    topic1 = @cluster.create_topic('test.kazoo.subscription.1', partitions: 1, replication_factor: 1)
+    topic2 = @cluster.create_topic('test.kazoo.subscription.2', partitions: 2, replication_factor: 1)
+    topic3 = @cluster.create_topic('test.kazoo.non_matching', partitions: 1, replication_factor: 1)
+
+    subscription = Kazoo::Subscription.build(/^test\.kazoo\.subscription\.\d+/)
+    assert_equal Set[topic1, topic2], Set.new(subscription.topics(@cluster))
+    assert_equal Set[topic1.partition(0), topic2.partition(0), topic2.partition(1)], Set.new(subscription.partitions(@cluster))
+
+    topic2.destroy
+
+    assert_equal Set[topic1], Set.new(subscription.topics(@cluster))
+    assert_equal Set[topic1.partition(0)], Set.new(subscription.partitions(@cluster))
+
+    topic1.destroy
+
+    assert_equal Set[], Set.new(subscription.topics(@cluster))
+    assert_equal Set[], Set.new(subscription.partitions(@cluster))
+
+    topic3.destroy
+  end
+end

data/test/functional/functional_topic_management_test.rb

@@ -4,6 +4,12 @@ class FunctionalTopicManagementTest < Minitest::Test
   def setup
     zookeeper = ENV["ZOOKEEPER_PEERS"] || "127.0.0.1:2181"
     @cluster = Kazoo.connect(zookeeper)
+
+    topic = @cluster.topic('test.kazoo')
+    topic.destroy if topic.exists?
+
+    topic = @cluster.topic('test.kazoo.config')
+    topic.destroy if topic.exists?
   end
 
   def test_create_and_delete_topic
@@ -14,16 +20,33 @@ class FunctionalTopicManagementTest < Minitest::Test
     assert_equal 8, topic.partitions.length
 
     topic.destroy
+
+    refute topic.exists?
+    refute @cluster.topics.key?(topic.name)
+  end
+
+  def test_adding_partitions_to_topic
+    topic = @cluster.create_topic('test.kazoo', partitions: 2, replication_factor: 1)
+
+    topic.add_partitions(partitions: 2, replication_factor: 1)
+
+    assert topic.partitions.all? { |partition| @cluster.brokers.values.include?(partition.leader) }
+    assert_equal 4, topic.partitions.length
+
     @cluster.reset_metadata
 
+    assert topic.partitions.all? { |partition| @cluster.brokers.values.include?(partition.leader) }
+    assert_equal 4, topic.partitions.length
+
+    topic.destroy
+
     refute topic.exists?
     refute @cluster.topics.key?(topic.name)
   end
 
   def test_topic_config_management
-    topic = @cluster.create_topic('test.kazoo.config', partitions: 1, replication_factor: 1)
+    topic = @cluster.create_topic('test.kazoo.config', partitions: 1, replication_factor: 1, config: { "flush.messages" => 1, "max.message.bytes" => 64000 })
 
-    topic.write_config("flush.messages" => 1, "max.message.bytes" => 64000)
     assert_equal "1", topic.config["flush.messages"]
     assert_equal "64000", topic.config["max.message.bytes"]
 

data/test/replica_assigner_test.rb

@@ -0,0 +1,80 @@
+require 'test_helper'
+
+class ReplicaAssignerTest < Minitest::Test
+  include MockCluster
+
+  def setup
+    @cluster = mock_cluster
+    @replica_assigner = Kazoo::ReplicaAssigner.new(@cluster)
+  end
+
+  def test_initial_counts
+    cluster_leader_count = @cluster.partitions.length
+    assert_equal cluster_leader_count, @replica_assigner.cluster_leader_count
+
+    assert_equal 2, @replica_assigner.broker_leaders[@cluster.brokers[1]]
+    assert_equal 2, @replica_assigner.broker_leaders[@cluster.brokers[2]]
+    assert_equal 1, @replica_assigner.broker_leaders[@cluster.brokers[3]]
+
+    cluster_replica_count = @cluster.partitions.inject(0) { |sum, p| sum + p.replicas.length }
+    assert_equal cluster_replica_count, @replica_assigner.cluster_replica_count
+
+    assert_equal 3, @replica_assigner.broker_replicas[@cluster.brokers[1]]
+    assert_equal 4, @replica_assigner.broker_replicas[@cluster.brokers[2]]
+    assert_equal 3, @replica_assigner.broker_replicas[@cluster.brokers[3]]
+  end
+
+  def test_assign_replicas_validation
+    assert_raises(Kazoo::ValidationError) { @replica_assigner.assign(4) }
+    assert_raises(Kazoo::ValidationError) { @replica_assigner.assign(0) }
+  end
+
+  def test_assign_replicas_deterministically
+    replicas_1 = @replica_assigner.assign(2)
+    assert_equal [@cluster.brokers[3], @cluster.brokers[1]], replicas_1
+
+    assert_equal 2, @replica_assigner.broker_leaders[@cluster.brokers[1]]
+    assert_equal 2, @replica_assigner.broker_leaders[@cluster.brokers[2]]
+    assert_equal 2, @replica_assigner.broker_leaders[@cluster.brokers[3]]
+
+    assert_equal 4, @replica_assigner.broker_replicas[@cluster.brokers[1]]
+    assert_equal 4, @replica_assigner.broker_replicas[@cluster.brokers[2]]
+    assert_equal 4, @replica_assigner.broker_replicas[@cluster.brokers[3]]
+
+    replicas_2 = @replica_assigner.assign(3)
+    assert_equal [@cluster.brokers[1], @cluster.brokers[2], @cluster.brokers[3]], replicas_2
+
+    assert_equal 3, @replica_assigner.broker_leaders[@cluster.brokers[1]]
+    assert_equal 2, @replica_assigner.broker_leaders[@cluster.brokers[2]]
+    assert_equal 2, @replica_assigner.broker_leaders[@cluster.brokers[3]]
+
+    assert_equal 5, @replica_assigner.broker_replicas[@cluster.brokers[1]]
+    assert_equal 5, @replica_assigner.broker_replicas[@cluster.brokers[2]]
+    assert_equal 5, @replica_assigner.broker_replicas[@cluster.brokers[3]]
+
+  end
+
+  def test_assign_replicas_evenly
+    initial_leaders = @replica_assigner.cluster_leader_count
+    initial_replicas = @replica_assigner.cluster_replica_count
+
+    1000.times do
+      replicas = @replica_assigner.assign(2)
+      assert_equal 2, replicas.length
+    end
+
+    # We expect the cluster-wide counters to be increased by the expected amount of replicas
+    final_leaders = @replica_assigner.cluster_leader_count
+    final_replicas = @replica_assigner.cluster_replica_count
+
+    assert_equal final_leaders, initial_leaders + 1000
+    assert_equal final_replicas, initial_replicas + 2000
+
+    # For the entire cluster, the number of leaders and replicas should be even for every broker.
+    leader_diff = @replica_assigner.broker_leaders.values.max - @replica_assigner.broker_leaders.values.min
+    assert leader_diff >= 0 && leader_diff <= 1
+
+    replica_diff = @replica_assigner.broker_replicas.values.max - @replica_assigner.broker_replicas.values.min
+    assert replica_diff >= 0 && replica_diff <= 1
+  end
+end

data/test/subscription_test.rb

@@ -0,0 +1,148 @@
+require 'test_helper'
+
+class SubscriptionTest < Minitest::Test
+  include MockCluster
+
+  def setup
+    @cluster = mock_cluster
+  end
+
+  def test_static_subscription_topics
+    subscription = Kazoo::Subscription.build(['test.1', 'nonexisting'])
+    topics = subscription.topics(@cluster)
+    assert_equal Set['test.1'], Set.new(topics.map(&:name))
+  end
+
+  def test_pattern_subscription_topics
+    subscription = Kazoo::Subscription.build(/^test\.\d+/, pattern: :white_list)
+    topics = subscription.topics(@cluster)
+    assert_equal Set['test.1', 'test.4'], Set.new(topics.map(&:name))
+
+    subscription = Kazoo::Subscription.build(/\.4/, pattern: :black_list)
+    topics = subscription.topics(@cluster)
+    assert_equal Set['test.1'], Set.new(topics.map(&:name))
+  end
+
+  def test_equality
+    subscription1 = Kazoo::Subscription.build(/^test\.\d+/, pattern: :white_list)
+    subscription2 = Kazoo::Subscription.build(/^test\.\d+/, pattern: :white_list)
+    subscription3 = Kazoo::Subscription.build(/^test\.\d+/, pattern: :black_list)
+    subscription4 = Kazoo::Subscription.build(/^test\.\d*/, pattern: :white_list)
+    assert subscription1 == subscription2
+    refute subscription1 == subscription3
+    refute subscription1 == subscription4
+
+    subscription1 = Kazoo::Subscription.build(:'test.1')
+    subscription2 = Kazoo::Subscription.build(['test.1'])
+    subscription3 = Kazoo::Subscription.build(['test.1', 'test.4'])
+    assert subscription1 == subscription2
+    refute subscription1 == subscription3
+  end
+
+  def test_subscription_from_json
+    timestamp_msec = 628232400123
+    timestamp = Time.at(BigDecimal.new(timestamp_msec) / BigDecimal.new(1000))
+
+    json_payload = JSON.generate(
+      version:      1,
+      timestamp:    timestamp_msec,
+      pattern:      "static",
+      subscription: { 'topic.1' => 1, 'topic.4' => 1 },
+    )
+
+    subscription = Kazoo::Subscription.from_json(json_payload)
+    assert_kind_of Kazoo::StaticSubscription, subscription
+    assert_equal timestamp, subscription.timestamp
+    assert_equal 1, subscription.version
+    assert_equal Set['topic.1', 'topic.4'], Set.new(subscription.topic_names)
+
+    json_payload = JSON.generate(
+      version:      1,
+      timestamp:    timestamp_msec,
+      pattern:      "black_list",
+      subscription: { "^test\\.\\d+" => 1 },
+    )
+
+    subscription = Kazoo::Subscription.from_json(json_payload)
+    assert_kind_of Kazoo::PatternSubscription, subscription
+    assert_equal timestamp, subscription.timestamp
+    assert_equal 1, subscription.version
+    assert subscription.black_list?
+    assert_equal %r{^test\.\d+}, subscription.regexp
+  end
+
+  def test_single_topic_static_subscription_json
+    subscription = Kazoo::Subscription.build('topic')
+    json = subscription.to_json
+
+    parsed_subscription = JSON.parse(json)
+    assert_equal 1, parsed_subscription.fetch('version')
+    assert_equal 'static', parsed_subscription.fetch('pattern')
+    assert_kind_of Integer, parsed_subscription.fetch('timestamp')
+
+    assert_kind_of Hash, parsed_subscription.fetch('subscription')
+    assert_equal 1, parsed_subscription.fetch('subscription').length
+    assert_equal 1, parsed_subscription.fetch('subscription').fetch('topic')
+  end
+
+  def test_multi_topic_static_subscription_json
+    subscription = Kazoo::Subscription.build([:topic1, :topic2])
+    json = subscription.to_json
+
+    parsed_subscription = JSON.parse(json)
+    assert_equal 1, parsed_subscription.fetch('version')
+    assert_equal 'static', parsed_subscription.fetch('pattern')
+    assert_kind_of Integer, parsed_subscription.fetch('timestamp')
+
+    assert_kind_of Hash, parsed_subscription.fetch('subscription')
+    assert_equal 2, parsed_subscription.fetch('subscription').length
+    assert_equal 1, parsed_subscription.fetch('subscription').fetch('topic1')
+    assert_equal 1, parsed_subscription.fetch('subscription').fetch('topic2')
+  end
+
+  def test_whitelist_subscription_json
+    subscription = Kazoo::Subscription.build(/^topic/)
+    json = subscription.to_json
+
+    parsed_subscription = JSON.parse(json)
+    assert_equal 1, parsed_subscription.fetch('version')
+    assert_equal 'white_list', parsed_subscription.fetch('pattern')
+    assert_kind_of Integer, parsed_subscription.fetch('timestamp')
+
+    assert_kind_of Hash, parsed_subscription.fetch('subscription')
+    assert_equal 1, parsed_subscription.fetch('subscription').length
+
+    assert_equal ["^topic"], parsed_subscription.fetch('subscription').keys
+    assert_equal [1], parsed_subscription.fetch('subscription').values
+  end
+
+  def test_blacklist_subscription_json
+    subscription = Kazoo::Subscription.build(/^topic/, pattern: :black_list)
+    json = subscription.to_json
+
+    parsed_subscription = JSON.parse(json)
+    assert_equal 1, parsed_subscription.fetch('version')
+    assert_equal 'black_list', parsed_subscription.fetch('pattern')
+    assert_kind_of Integer, parsed_subscription.fetch('timestamp')
+
+    assert_kind_of Hash, parsed_subscription.fetch('subscription')
+    assert_equal 1, parsed_subscription.fetch('subscription').length
+
+    assert_equal ["^topic"], parsed_subscription.fetch('subscription').keys
+    assert_equal [1], parsed_subscription.fetch('subscription').values
+  end
+
+  def test_subscription_has_topic
+    static_subscription = Kazoo::Subscription.build(%w[topic.1 topic.4], pattern: :black_list)
+    assert static_subscription.has_topic?(@cluster.topic('topic.1'))
+    refute static_subscription.has_topic?(@cluster.topic('test.topic'))
+
+    whitelist_subscription = Kazoo::Subscription.build(/^topic\./, pattern: :white_list)
+    assert whitelist_subscription.has_topic?(@cluster.topic('topic.1'))
+    refute whitelist_subscription.has_topic?(@cluster.topic('test.topic'))
+
+    blacklist_subscription = Kazoo::Subscription.build(/^topic\./, pattern: :black_list)
+    refute blacklist_subscription.has_topic?(@cluster.topic('topic.1'))
+    assert blacklist_subscription.has_topic?(@cluster.topic('test.topic'))
+  end
+end