deimos-ruby 2.4.0.pre.beta8 → 2.4.0.pre.beta10

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c1b2dfa779652664bad07fb938a1f4472bca69b7c61c21b4cb6fe028c86061a9
-  data.tar.gz: 0bf9e789a83563e10b3bd885de8736c5b4487268e4b9cc1225ac8f8f3ee2a866
+  metadata.gz: 2bfb935df5a67962364fbb6774c4dd03e2e56cbdb4d6dad7d15cb034b9522dee
+  data.tar.gz: c55979be4de60e902b02f80c48bce656d550d89fdb712938462d1f8b1af1f38b
 SHA512:
-  metadata.gz: b089bcb150fabaabf23b3d4431f8de937080961b298c3160c7b02db5e41468f524b9d0f8b634f3b69ae99dadf12e2c67b06d420c0821f24ff42eb4c7f6222b2e
-  data.tar.gz: 41bf155c6813975fdda7818d9199a1334572e70601e7a16076430325f0e1bbe4a9983c0169eaac0e91ede2081947c762a1e6a618f379ad1ade21c78711cfcd49
+  metadata.gz: 9fc0b3ed398049f73002d77c8b103f69b1b569616a9bfa960921edc3947702fc3affdc7d39676db43e8e935df2e08b2bb97726ea93060503d3f53f00a36a79bc
+  data.tar.gz: cae1f740a93966da3076095c1adbb1e9fd125ee4531947d15913d98795d460c21d3a98a6840137ef40a1d8d13374833065a59ee6e04196cb58e1c9b3f4acc1b8

data/README.md

@@ -281,9 +281,15 @@ MyProducer.publish({
                    })
 ```
 
+### Protobuf and Key Schemas
+
 > [!IMPORTANT]
 > Protobuf should *not* be used as a key schema, since the binary encoding is [unstable](https://protobuf.dev/programming-guides/encoding/#implications) and may break partitioning. Deimos will automatically convert keys to sorted JSON, and will use JSON Schema in the schema registry.
 
+To enable integration with [buf](https://buf.build/), Deimos provides a Rake task that will auto-generate `.proto` files. This task can be run in CI before running `buf push`. This way, downstream systems that want to use the generated key schemas can do so using their own tech stack.
+
+    bundle exec rake deimos:generate_key_protos
+
 ## Instrumentation
 
 Deimos will send events through the [Karafka instrumentation monitor](https://karafka.io/docs/Monitoring-and-Logging/#subscribing-to-the-instrumentation-events). 

data/lib/deimos/config/configuration.rb

@@ -190,6 +190,9 @@ module Deimos
       # would replace a prefixed with the given key with the module name SchemaClasses.
       # @return [Hash]
       setting :schema_namespace_map, {}
+
+      # The base directory for generated protobuf key schemas.
+      setting :proto_schema_key_path, 'protos'
     end
 
     # The configured metrics provider.

data/lib/deimos/ext/producer_route.rb

@@ -10,9 +10,10 @@ module Deimos
       end
 
       def producer_class
-        self.producer_classes.first
+        self.producer_classes&.first
       end
     end
+
     module Topic
       (FIELDS + [:producer_class]).each do |field|
         define_method(field) do |*args|

data/lib/deimos/schema_backends/base.rb

@@ -41,6 +41,11 @@ module Deimos
         @namespace = namespace
       end
 
+      # @return [String]
+      def inspect
+        "Type #{self.class.name.demodulize} Schema: #{@namespace}.#{@schema} Key schema: #{@key_schema}"
+      end
+
       # @return [Boolean]
       def supports_key_schemas?
         false

data/lib/deimos/schema_backends/proto_schema_registry.rb

@@ -56,6 +56,29 @@ module Deimos
         )
       end
 
+      # @param file [String]
+      # @param field_name [String]
+      def write_key_proto(file, field_name)
+        return if field_name.nil?
+
+        proto = proto_schema
+        package = proto.file_descriptor.to_proto.package
+        writer = SchemaRegistry::Output::ProtoText::Writer.new
+        info = SchemaRegistry::Output::ProtoText::ParseInfo.new(writer, package)
+        writer.write_line(%(syntax = "proto3";))
+        writer.write_line("package #{package};")
+        writer.writenl
+
+        field = proto.to_proto.field.find { |f| f.name == field_name.to_s }
+        writer.write_line("message #{proto.to_proto.name}Key {")
+        writer.indent
+        SchemaRegistry::Output::ProtoText.write_field(info, field)
+        writer.dedent
+        writer.write_line('}')
+        path = "#{file}/#{package.gsub('.', '/')}/#{proto.to_proto.name.underscore}_key.proto"
+        File.write(path, writer.string)
+      end
+
     end
   end
 end

data/lib/deimos/test_helpers.rb

@@ -58,6 +58,7 @@ module Deimos
           config.deserializers[:key].try(:reset_backend)
         end
         yield
+      ensure
         Deimos.mock_backends = false
       end
 

data/lib/deimos/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Deimos
-  VERSION = '2.4.0-beta8'
+  VERSION = '2.4.0-beta10'
 end

data/lib/deimos.rb

@@ -72,9 +72,18 @@ module Deimos
       klass
     end
 
+    # @param topic_name [String]
+    # @return [SchemaBackends::Base]
+    def schema_backend_for(topic_name)
+      config = Deimos.karafka_config_for(topic: topic_name)
+      self.schema_backend(schema: config.schema,
+                          namespace: config.namespace,
+                          backend: config.schema_backend || Deimos.config.schema.backend)
+    end
+
     # @param schema [String, Symbol]
     # @param namespace [String]
-    # @return [Class<Deimos::SchemaBackends::Base>]
+    # @return [Deimos::SchemaBackends::Base]
     def schema_backend(schema:, namespace:, backend: Deimos.config.schema.backend)
       if config.schema.use_schema_classes
         # Initialize an instance of the provided schema

data/lib/tasks/deimos.rake

@@ -2,6 +2,7 @@
 
 require 'generators/deimos/schema_class_generator'
 require 'optparse'
+require 'deimos/schema_backends/proto_schema_registry'
 
 namespace :deimos do
   desc 'Starts Deimos in the rails environment'
@@ -40,4 +41,17 @@ namespace :deimos do
     Deimos::Generators::SchemaClassGenerator.start
   end
 
+  desc 'Output Protobuf key schemas'
+  task generate_key_protos: :environment do
+    puts "Generating Protobuf key schemas"
+    Deimos.karafka_configs.each do |config|
+      next if config.try(:producer_class).blank? ||
+              !Deimos.schema_backend_for(config.name).is_a?(Deimos::SchemaBackends::ProtoSchemaRegistry) ||
+              config.deserializers[:key].try(:key_field).blank?
+
+      puts "Writing key proto for #{config.name}"
+      backend = Deimos.schema_backend_for(config.name)
+      backend.write_key_proto(Deimos.config.schema.proto_schema_key_path, config.deserializers[:key].key_field)
+    end
+  end
 end

data/spec/deimos_spec.rb

@@ -82,4 +82,74 @@ describe Deimos do
     expect(producer3.config.kafka[:'bootstrap.servers']).to eq('broker2:9092')
   end
 
+  describe '#schema_backend_for' do
+    it 'should return a schema backend for a given topic' do
+      Karafka::App.routes.redraw do
+        topic 'backend-test-topic' do
+          active false
+          schema 'MySchema'
+          namespace 'com.my-namespace'
+          key_config none: true
+        end
+      end
+
+      backend = described_class.schema_backend_for('backend-test-topic')
+      expect(backend.inspect).to eq('Type AvroValidation Schema: com.my-namespace.MySchema Key schema: ')
+    end
+
+    it 'should use topic-specific schema_backend if configured' do
+      Karafka::App.routes.redraw do
+        topic 'backend-test-topic-local' do
+          active false
+          schema 'MySchema'
+          namespace 'com.my-namespace'
+          key_config none: true
+          schema_backend :avro_local
+        end
+      end
+
+      backend = described_class.schema_backend_for('backend-test-topic-local')
+      expect(backend).to be_a(Deimos::SchemaBackends::AvroLocal)
+    end
+
+    it 'should fall back to global schema backend config' do
+      described_class.config.schema.backend = :avro_validation
+      Karafka::App.routes.redraw do
+        topic 'backend-test-topic-global' do
+          active false
+          schema 'MySchema'
+          namespace 'com.my-namespace'
+          key_config none: true
+        end
+      end
+
+      backend = described_class.schema_backend_for('backend-test-topic-global')
+      expect(backend).to be_a(Deimos::SchemaBackends::AvroValidation)
+    end
+  end
+
+  describe '#schema_backend_class with mock_backends' do
+    after(:each) do
+      described_class.mock_backends = false
+    end
+
+    it 'should return the mock backend when mock_backends is true' do
+      described_class.mock_backends = true
+      klass = described_class.schema_backend_class(backend: :avro_schema_registry)
+      expect(klass).to eq(Deimos::SchemaBackends::AvroValidation)
+    end
+
+    it 'should return the real backend when mock_backends is false' do
+      described_class.mock_backends = false
+      klass = described_class.schema_backend_class(backend: :avro_schema_registry)
+      expect(klass).to eq(Deimos::SchemaBackends::AvroSchemaRegistry)
+    end
+
+    it 'should use proto_local for proto_schema_registry when mocked' do
+      described_class.mock_backends = true
+      klass = described_class.schema_backend_class(backend: :proto_schema_registry)
+      expect(klass).to eq(Deimos::SchemaBackends::ProtoLocal)
+    end
+  end
+
 end

data/spec/gen/sample/v1/sample_pb.rb

@@ -1,5 +1,4 @@
 # frozen_string_literal: true
-
 # Generated by the protocol buffer compiler.  DO NOT EDIT!
 # source: sample/v1/sample.proto
 
@@ -7,14 +6,15 @@ require 'google/protobuf'
 
 require 'google/protobuf/timestamp_pb'
 
-descriptor_data = "\n\x16sample/v1/sample.proto\x12\tsample.v1\x1a\x1fgoogle/protobuf/timestamp.proto\"7\n\rNestedMessage\x12\x12\n\nnested_str\x18\x01 \x01(\t\x12\x12\n\nnested_num\x18\x02 \x01(\x05\"\xdb\x02\n\rSampleMessage\x12\x0b\n\x03str\x18\x01 \x01(\t\x12\x0b\n\x03num\x18\x02 \x01(\x05\x12\x0f\n\x07str_arr\x18\x03 \x03(\t\x12\x0c\n\x04\x66lag\x18\x04 \x01(\x08\x12-\n\ttimestamp\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12*\n\x06nested\x18\x06 \x01(\x0b\x32\x18.sample.v1.NestedMessageH\x00\x12\x13\n\tunion_str\x18\x07 \x01(\tH\x00\x12\x32\n\x10non_union_nested\x18\x08 \x01(\x0b\x32\x18.sample.v1.NestedMessage\x12\x35\n\x07str_map\x18\t \x03(\x0b\x32$.sample.v1.SampleMessage.StrMapEntry\x1a-\n\x0bStrMapEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x42\x07\n\x05unionb\x06proto3"
 
-pool = Google::Protobuf::DescriptorPool.generated_pool
+descriptor_data = "\n\x16sample/v1/sample.proto\x12\tsample.v1\x1a\x1fgoogle/protobuf/timestamp.proto\"7\n\rNestedMessage\x12\x12\n\nnested_str\x18\x01 \x01(\t\x12\x12\n\nnested_num\x18\x02 \x01(\x05\"\xef\x02\n\rSampleMessage\x12\x0b\n\x03str\x18\x01 \x01(\t\x12\x0b\n\x03num\x18\x02 \x01(\x05\x12\x0f\n\x07str_arr\x18\x03 \x03(\t\x12\x0c\n\x04\x66lag\x18\x04 \x01(\x08\x12-\n\ttimestamp\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12*\n\x06nested\x18\x06 \x01(\x0b\x32\x18.sample.v1.NestedMessageH\x00\x12\x13\n\tunion_str\x18\x07 \x01(\tH\x00\x12\x32\n\x10non_union_nested\x18\x08 \x01(\x0b\x32\x18.sample.v1.NestedMessage\x12\x35\n\x07str_map\x18\t \x03(\x0b\x32$.sample.v1.SampleMessage.StrMapEntry\x12\x12\n\nmessage_id\x18\n \x01(\t\x1a-\n\x0bStrMapEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x42\x07\n\x05unionb\x06proto3"
+
+pool = ::Google::Protobuf::DescriptorPool.generated_pool
 pool.add_serialized_file(descriptor_data)
 
 module Sample
   module V1
-    NestedMessage = ::Google::Protobuf::DescriptorPool.generated_pool.lookup('sample.v1.NestedMessage').msgclass
-    SampleMessage = ::Google::Protobuf::DescriptorPool.generated_pool.lookup('sample.v1.SampleMessage').msgclass
+    NestedMessage = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("sample.v1.NestedMessage").msgclass
+    SampleMessage = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("sample.v1.SampleMessage").msgclass
   end
 end

data/spec/message_spec.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative 'gen/sample/v1/sample_pb'
+
 RSpec.describe(Deimos::Message) do
   it 'should detect tombstones' do
     expect(described_class.new(nil, key: 'key1')).
@@ -36,4 +38,131 @@ RSpec.describe(Deimos::Message) do
         to include(headers: { a: 1 })
     end
   end
+
+  describe '#add_fields' do
+    context 'with Hash payloads' do
+      it 'should add message_id if not present' do
+        message = described_class.new({ some_field: 'value' }, key: 'key1')
+        message.add_fields(%w(message_id timestamp))
+        expect(message.payload['message_id']).to be_present
+        expect(message.payload['message_id']).to match(/\A[0-9a-f-]{36}\z/)
+      end
+
+      it 'should not overwrite existing message_id' do
+        message = described_class.new({ some_field: 'value', message_id: 'existing-id' }, key: 'key1')
+        message.add_fields(%w(message_id timestamp))
+        expect(message.payload['message_id']).to eq('existing-id')
+      end
+
+      it 'should add timestamp if not present' do
+        freeze_time = Time.zone.local(2025, 1, 1, 12, 0, 0)
+        travel_to(freeze_time) do
+          message = described_class.new({ some_field: 'value' }, key: 'key1')
+          message.add_fields(%w(message_id timestamp))
+          expect(message.payload['timestamp']).to be_present
+          expect(message.payload['timestamp']).to include('2025')
+        end
+      end
+
+      it 'should not overwrite existing timestamp' do
+        message = described_class.new({ some_field: 'value', timestamp: '2020-01-01' }, key: 'key1')
+        message.add_fields(%w(message_id timestamp))
+        expect(message.payload['timestamp']).to eq('2020-01-01')
+      end
+
+      it 'should not add fields that are not in the list' do
+        message = described_class.new({ some_field: 'value' }, key: 'key1')
+        message.add_fields(%w(other_field))
+        expect(message.payload['message_id']).to be_nil
+        expect(message.payload['timestamp']).to be_nil
+      end
+
+      it 'should do nothing for empty payload' do
+        message = described_class.new({}, key: 'key1')
+        message.add_fields(%w(message_id timestamp))
+        expect(message.payload['message_id']).to be_nil
+        expect(message.payload['timestamp']).to be_nil
+      end
+
+      it 'should do nothing for nil payload' do
+        message = described_class.new(nil, key: 'key1')
+        message.add_fields(%w(message_id timestamp))
+        expect(message.payload).to be_nil
+      end
+    end
+
+    context 'with Protobuf payloads' do
+      it 'should add message_id to protobuf objects if not present' do
+        proto_payload = Sample::V1::SampleMessage.new(str: 'test', num: 123)
+        expect(proto_payload.message_id).to eq('')
+
+        message = described_class.new(proto_payload, key: 'key1')
+        message.add_fields(%w(message_id))
+
+        expect(proto_payload.message_id).to match(/\A[0-9a-f-]{36}\z/)
+      end
+
+      it 'should not overwrite existing message_id on protobuf objects' do
+        proto_payload = Sample::V1::SampleMessage.new(str: 'test', num: 123, message_id: 'existing-id')
+
+        message = described_class.new(proto_payload, key: 'key1')
+        message.add_fields(%w(message_id))
+
+        expect(proto_payload.message_id).to eq('existing-id')
+      end
+
+      it 'should add timestamp to protobuf objects if not present' do
+        proto_payload = Sample::V1::SampleMessage.new(str: 'test', num: 123)
+        expect(proto_payload.timestamp).to be_nil
+
+        message = described_class.new(proto_payload, key: 'key1')
+        message.add_fields(%w(timestamp))
+
+        expect(proto_payload.timestamp).to be_a(Google::Protobuf::Timestamp)
+        expect(proto_payload.timestamp.seconds).to be > 0
+      end
+
+      it 'should not overwrite existing timestamp on protobuf objects' do
+        existing_time = Google::Protobuf::Timestamp.new(seconds: 1_577_836_800) # 2020-01-01
+        proto_payload = Sample::V1::SampleMessage.new(str: 'test', num: 123, timestamp: existing_time)
+
+        message = described_class.new(proto_payload, key: 'key1')
+        message.add_fields(%w(timestamp))
+
+        expect(proto_payload.timestamp.seconds).to eq(1_577_836_800)
+      end
+
+      it 'should add both message_id and timestamp to protobuf objects' do
+        proto_payload = Sample::V1::SampleMessage.new(str: 'test', num: 123)
+
+        message = described_class.new(proto_payload, key: 'key1')
+        message.add_fields(%w(message_id timestamp))
+
+        expect(proto_payload.message_id).to match(/\A[0-9a-f-]{36}\z/)
+        expect(proto_payload.timestamp).to be_a(Google::Protobuf::Timestamp)
+      end
+
+      it 'should not modify protobuf when field not in list' do
+        proto_payload = Sample::V1::SampleMessage.new(str: 'test', num: 123)
+
+        message = described_class.new(proto_payload, key: 'key1')
+        message.add_fields(%w(other_field))
+
+        expect(proto_payload.message_id).to eq('')
+        expect(proto_payload.timestamp).to be_nil
+      end
+
+      it 'should handle protobuf with nested messages' do
+        nested = Sample::V1::NestedMessage.new(nested_str: 'nested', nested_num: 456)
+        proto_payload = Sample::V1::SampleMessage.new(str: 'test', num: 123, non_union_nested: nested)
+
+        message = described_class.new(proto_payload, key: 'key1')
+        message.add_fields(%w(message_id timestamp))
+
+        expect(proto_payload.message_id).to match(/\A[0-9a-f-]{36}\z/)
+        expect(proto_payload.timestamp).to be_a(Google::Protobuf::Timestamp)
+        expect(proto_payload.non_union_nested.nested_str).to eq('nested')
+      end
+    end
+  end
 end

data/spec/protos/sample/v1/sample.proto

@@ -21,4 +21,5 @@ message SampleMessage {
   }
   NestedMessage non_union_nested = 8;
   map<string, string> str_map = 9;
+  string message_id = 10;
 }

data/spec/rake_spec.rb

@@ -2,6 +2,12 @@
 
 require 'rake'
 require 'rails'
+require 'deimos/schema_backends/proto_schema_registry'
+require 'deimos/schema_backends/avro_validation'
+require_relative 'gen/sample/v1/sample_pb'
+require 'fileutils'
+require 'tmpdir'
+
 Rails.logger = Logger.new(STDOUT)
 load("#{__dir__}/../lib/tasks/deimos.rake")
 
@@ -14,4 +20,171 @@ describe 'Rakefile' do
     expect(Karafka::Server).to receive(:run)
     Rake::Task['deimos:start'].invoke
   end
+
+  describe 'generate_key_protos' do
+    let(:temp_dir) { Dir.mktmpdir('deimos_rake_test') }
+
+    before(:each) do
+      Rake::Task['deimos:generate_key_protos'].reenable
+      # Create the package directory structure
+      FileUtils.mkdir_p(File.join(temp_dir, 'sample/v1'))
+      # Configure Deimos to use the temp directory
+      Deimos.config.schema.proto_schema_key_path = temp_dir
+    end
+
+    after(:each) do
+      FileUtils.rm_rf(temp_dir)
+    end
+
+    it 'should skip configs without producer_class' do
+      Karafka::App.routes.redraw do
+        topic 'no-producer-topic' do
+          active false
+          schema 'sample.v1.SampleMessage'
+          namespace 'sample.v1'
+          key_config field: 'str'
+          schema_backend :proto_schema_registry
+        end
+      end
+
+      Rake::Task['deimos:generate_key_protos'].invoke
+
+      # No files should be written since there's no producer_class
+      expect(Dir.glob("#{temp_dir}/**/*.proto")).to be_empty
+    end
+
+    it 'should skip non-protobuf backends' do
+      producer_class = Class.new(Deimos::Producer)
+      stub_const('AvroProducer', producer_class)
+
+      Karafka::App.routes.redraw do
+        topic 'avro-topic' do
+          active false
+          schema 'MySchema'
+          namespace 'com.my-namespace'
+          key_config field: 'test_id'
+          producer_class AvroProducer
+        end
+      end
+
+      Rake::Task['deimos:generate_key_protos'].invoke
+
+      # No files should be written for non-protobuf backends
+      expect(Dir.glob("#{temp_dir}/**/*.proto")).to be_empty
+    end
+
+    it 'should skip configs without key_field' do
+      producer_class = Class.new(Deimos::Producer)
+      stub_const('NoKeyProducer', producer_class)
+
+      Karafka::App.routes.redraw do
+        topic 'no-key-topic' do
+          active false
+          schema 'sample.v1.SampleMessage'
+          namespace 'sample.v1'
+          key_config none: true
+          producer_class NoKeyProducer
+          schema_backend :proto_schema_registry
+        end
+      end
+
+      Rake::Task['deimos:generate_key_protos'].invoke
+
+      # No files should be written since there's no key_field
+      expect(Dir.glob("#{temp_dir}/**/*.proto")).to be_empty
+    end
+
+    it 'should write key proto for protobuf topics with string key_field' do
+      producer_class = Class.new(Deimos::Producer)
+      stub_const('ProtoProducer', producer_class)
+
+      Karafka::App.routes.redraw do
+        topic 'proto-topic' do
+          active false
+          schema 'sample.v1.SampleMessage'
+          namespace 'sample.v1'
+          key_config field: 'str'
+          producer_class ProtoProducer
+          schema_backend :proto_schema_registry
+        end
+      end
+
+      Rake::Task['deimos:generate_key_protos'].invoke
+
+      # Verify the file was created with correct content
+      proto_file = File.join(temp_dir, 'sample/v1/sample_message_key.proto')
+      expect(File).to exist(proto_file)
+
+      content = File.read(proto_file)
+      expect(content).to include('syntax = "proto3";')
+      expect(content).to include('package sample.v1;')
+      expect(content).to include('message SampleMessageKey {')
+      expect(content).to include('string str = 1;')
+    end
+
+    it 'should write key proto for protobuf topics with int32 key_field' do
+      producer_class = Class.new(Deimos::Producer)
+      stub_const('ProtoProducer', producer_class)
+
+      Karafka::App.routes.redraw do
+        topic 'proto-topic-int' do
+          active false
+          schema 'sample.v1.SampleMessage'
+          namespace 'sample.v1'
+          key_config field: 'num'
+          producer_class ProtoProducer
+          schema_backend :proto_schema_registry
+        end
+      end
+
+      Rake::Task['deimos:generate_key_protos'].invoke
+
+      proto_file = File.join(temp_dir, 'sample/v1/sample_message_key.proto')
+      expect(File).to exist(proto_file)
+
+      content = File.read(proto_file)
+      expect(content).to include('syntax = "proto3";')
+      expect(content).to include('package sample.v1;')
+      expect(content).to include('message SampleMessageKey {')
+      expect(content).to include('int32 num = 2;')
+    end
+
+    it 'should process multiple topics' do
+      producer_class1 = Class.new(Deimos::Producer)
+      producer_class2 = Class.new(Deimos::Producer)
+      stub_const('ProtoProducer1', producer_class1)
+      stub_const('ProtoProducer2', producer_class2)
+
+      Karafka::App.routes.redraw do
+        topic 'proto-topic-1' do
+          active false
+          schema 'sample.v1.SampleMessage'
+          namespace 'sample.v1'
+          key_config field: 'str'
+          producer_class ProtoProducer1
+          schema_backend :proto_schema_registry
+        end
+        topic 'proto-topic-2' do
+          active false
+          schema 'sample.v1.SampleMessage'
+          namespace 'sample.v1'
+          key_config field: 'num'
+          producer_class ProtoProducer2
+          schema_backend :proto_schema_registry
+        end
+      end
+
+      Rake::Task['deimos:generate_key_protos'].invoke
+
+      # Both topics use the same schema, so only one file is created
+      # but it should be called twice (once per topic)
+      proto_file = File.join(temp_dir, 'sample/v1/sample_message_key.proto')
+      expect(File).to exist(proto_file)
+
+      # The second topic overwrites the first, so we should see the 'num' field
+      content = File.read(proto_file)
+      expect(content).to include('message SampleMessageKey {')
+      expect(content).to include('int32 num = 2;')
+    end
+  end
 end

data/spec/schema_backends/proto_schema_registry_spec.rb

@@ -247,4 +247,10 @@ RSpec.describe Deimos::SchemaBackends::ProtoSchemaRegistry do
       expect(fields.map(&:name)).to include('str', 'num', 'str_arr', 'flag', 'timestamp')
     end
   end
+
+  describe '.mock_backend' do
+    it 'should return :proto_local' do
+      expect(described_class.mock_backend).to eq(:proto_local)
+    end
+  end
 end

data/spec/test_helpers_spec.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+require 'deimos/test_helpers'
+
+RSpec.describe Deimos::TestHelpers do
+
+  describe '.with_mock_backends' do
+    before(:each) do
+      Karafka::App.routes.redraw do
+        topic 'mock-backend-test' do
+          active false
+          schema 'MySchema'
+          namespace 'com.my-namespace'
+          key_config field: 'test_id'
+        end
+      end
+    end
+
+    after(:each) do
+      Deimos.mock_backends = false
+    end
+
+    it 'should set mock_backends to true during the block' do
+      expect(Deimos.mock_backends).to be_falsey
+
+      described_class.with_mock_backends do
+        expect(Deimos.mock_backends).to be(true)
+      end
+
+      expect(Deimos.mock_backends).to be(false)
+    end
+
+    it 'should reset mock_backends after the block completes' do
+      described_class.with_mock_backends do
+        # noop
+      end
+      expect(Deimos.mock_backends).to be(false)
+    end
+
+    it 'should reset mock_backends even if an error occurs in the block' do
+      expect {
+        described_class.with_mock_backends do
+          raise 'Test error'
+        end
+      }.to raise_error('Test error')
+      expect(Deimos.mock_backends).to be(false)
+    end
+
+    it 'should reset backends on deserializers' do
+      config = Deimos.karafka_config_for(topic: 'mock-backend-test')
+      payload_transcoder = config.deserializers[:payload]
+      key_transcoder = config.deserializers[:key]
+
+      expect(payload_transcoder).to receive(:reset_backend)
+      expect(key_transcoder).to receive(:reset_backend) if key_transcoder.respond_to?(:reset_backend)
+
+      described_class.with_mock_backends do
+        # noop
+      end
+    end
+
+    it 'should use mock backends for schema_backend_class' do
+      described_class.with_mock_backends do
+        # avro_schema_registry should return avro_validation when mocked
+        klass = Deimos.schema_backend_class(backend: :avro_schema_registry)
+        expect(klass).to eq(Deimos::SchemaBackends::AvroValidation)
+      end
+
+      # Outside the block, should return the real class
+      klass = Deimos.schema_backend_class(backend: :avro_schema_registry)
+      expect(klass).to eq(Deimos::SchemaBackends::AvroSchemaRegistry)
+    end
+  end
+
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: deimos-ruby
 version: !ruby/object:Gem::Version
-  version: 2.4.0.pre.beta8
+  version: 2.4.0.pre.beta10
 platform: ruby
 authors:
 - Daniel Orner
@@ -669,6 +669,7 @@ files:
 - spec/snapshots/producers_with_key-no-nest.snap
 - spec/snapshots/producers_with_key.snap
 - spec/spec_helper.rb
+- spec/test_helpers_spec.rb
 - spec/utils/db_poller_spec.rb
 - spec/utils/deadlock_retry_spec.rb
 - spec/utils/outbox_producer_spec.rb