gRPCのレスポンスをファイルに記録してrspecで再生できるようにしたやつ

grpc_vcr.rb

require "grpc"

module GRPC::VCR
  GRPC_VCR_FiXTURE_PATH = SPEC_PATH.join("fixtures/grpc_vcr")

  def grpc_interceptor(example)
    cassette = example.metadata[:grpc_vcr]
    return unless cassette&.is_a? GRPC::VCR::Cassette

    cassette.exist? ?
      GRPC::VCR::PlayInterceptor.new(cassette) :
      GRPC::VCR::RecordInterceptor.new(cassette)
  end
end

module GRPC::VCR::Metadata
  def self.configure(config)
    when_tagged_with_grpc_vcr = { grpc_vcr: lambda {|v| !!v } }

    config.before(:each, when_tagged_with_grpc_vcr) do |ex|
      example = ex.respond_to?(:metadata) ? ex : ex.example

      name = example.metadata[:grpc_vcr]
      example.metadata[:grpc_vcr] = GRPC::VCR::Cassette.new(name)
    end

    config.after(:each, when_tagged_with_grpc_vcr) do |ex|
      example = ex.respond_to?(:metadata) ? ex : ex.example
      cassette = example.metadata[:grpc_vcr]
      cassette.save if !cassette.exist? && !ex.exception
    end
  end
end

class GRPC::VCR::Cassette
  attr_accessor :name, :path

  def initialize(name)
    self.name = name
    self.path = GRPC::VCR::GRPC_VCR_FiXTURE_PATH.join(name + ".json")

    @exist = File.exist?(path)
  end

  def exist?
    @exist
  end

  def append_record(res)
    @records ||= []
    record = res.as_json
    record["message"] = Base64.encode64(record["message"]) if record["message"]
    @records << record.as_json
  end

  def load
    @loaded = true
    @records = JSON.parse(File.open(path).read)
  end

  def save
    FileUtils.mkdir_p(File.dirname(path))

    File.open(path, "w"){|f| f.write @records.to_json }
  end

  def next_record
    self.load unless @loaded
    @iter ||= @records.each
    record = @iter.next rescue StopIteration
    return unless record # TODO: raise error

    record["message"] = Base64.decode64(record["message"]) if record["message"]
    record
  end
end

class GRPC::VCR::Interceptor < GRPC::ClientInterceptor
  attr_accessor :cassette

  def initialize(cassette)
    self.cassette = cassette
  end

  def extract_internal_call(call)
    call.instance_variable_get("@wrapped")&.instance_variable_get("@call")
  end

  def request_response(request:, call:, method:, metadata:)
    hook_run_batch_on_call(call, :request_response)
    yield
  end

  def client_streamer(call:, method:, metadata:)
    hook_run_batch_on_call(call, :client_streamer)
    yield
  end

  def server_streamer(request:, call:, method:, metadata:)
    hook_run_batch_on_call(call, :server_streamer)
    yield
  end

  def bidi_streamer(requests:, call:, method:, metadata:)
    hook_run_batch_on_call(call, :bidi_streamer)
    yield
  end
end

class GRPC::VCR::PlayInterceptor < GRPC::VCR::Interceptor
  def hook_run_batch_on_call(call, from)
    internal_call = extract_internal_call(call)
    internal_call.instance_variable_set("@grpc_vcr_cassette", self.cassette)

    def internal_call.run_batch(opts)
      record = @grpc_vcr_cassette.next_record

      status = nil
      status_attributes = record["status"]
      if status_attributes
        status = Struct::Status.new(
          status_attributes["code"],
          status_attributes["details"],
          status_attributes["metadata"],
          status_attributes["debug_error_string"]
        )
      end

      # puts "GRPC::VCR::PlayInterceptor: run_batch : play #{opts} => #{record["message"]&.size.to_i}"

      Struct::BatchResult.new(
        record["send_message"],
        record["send_metadata"],
        record["send_close"],
        record["send_status"],
        record["message"],
        record["metadata"],
        status,
        record["cancelled"]
      )
    end
  end
end

class GRPC::VCR::RecordInterceptor < GRPC::VCR::Interceptor
  def hook_run_batch_on_call(call, from)
    internal_call = extract_internal_call(call)
    internal_call.instance_variable_set("@grpc_vcr_cassette", self.cassette)

    def internal_call.run_batch(opts)
      res = super

      # puts "GRPC::VCR::RecordInterceptor: run_batch : saved #{opts} => #{res.message&.size.to_i}"
      @grpc_vcr_cassette.append_record(res)

      res
    end
  end
end