chore(ruby): Add mock server tests (#600)

Author: aakashanandg

spannerlib/wrappers/spannerlib-ruby/.rubocop.yml

@@ -30,3 +30,25 @@ RSpec/DescribeClass:
 
 Style/StringLiterals:
   EnforcedStyle: double_quotes
+
+Style/FormatStringToken:
+  EnforcedStyle: template
+Metrics/ClassLength:
+  Exclude:
+    - 'spec/mock_server/**/*'
+
+Metrics/MethodLength:
+  Exclude:
+    - 'spec/mock_server/**/*'
+
+Metrics/AbcSize:
+  Exclude:
+    - 'spec/mock_server/**/*'
+
+Metrics/CyclomaticComplexity:
+  Exclude:
+    - 'spec/mock_server/**/*'
+
+Metrics/PerceivedComplexity:
+  Exclude:
+    - 'spec/mock_server/**/*'

spannerlib/wrappers/spannerlib-ruby/Gemfile

@@ -7,6 +7,7 @@ gemspec
 gem "rake", "~> 13.0"
 
 group :development, :test do
+  gem "minitest", "~> 5.0"
   gem "rake-compiler", "~> 1.0"
   gem "rspec", "~> 3.0"
   gem "rubocop", require: false

spannerlib/wrappers/spannerlib-ruby/spec/mock_server/spanner_mock_server.rb

@@ -0,0 +1,283 @@
+# frozen_string_literal: true
+
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require_relative "statement_result"
+require "google/rpc/error_details_pb"
+require "google/spanner/v1/spanner_pb"
+require "google/spanner/v1/spanner_services_pb"
+require "google/cloud/spanner/v1/spanner"
+
+require "grpc"
+require "gapic/grpc/service_stub"
+require "securerandom"
+
+# Mock implementation of Spanner
+class SpannerMockServer < Google::Cloud::Spanner::V1::Spanner::Service
+  attr_reader :requests
+
+  def initialize
+    super
+    @statement_results = {}
+    @sessions = {}
+    @transactions = {}
+    @aborted_transactions = {}
+    @requests = []
+    @errors = {}
+    @session_counter = 0
+
+    put_statement_result "SELECT 1", StatementResult.create_select1_result
+    put_statement_result "INSERT INTO test_table (id, name) VALUES ('1', 'Alice')", StatementResult.create_update_count_result(1)
+
+    dialect_sql = "select option_value from information_schema.database_options where option_name='database_dialect'"
+    dialect_result = StatementResult.create_dialect_result
+
+    put_statement_result dialect_sql, dialect_result
+  end
+
+  def put_statement_result(sql, result)
+    @statement_results[sql] = result
+  end
+
+  def push_error(sql_or_method, error)
+    @errors[sql_or_method] = [] unless @errors[sql_or_method]
+    @errors[sql_or_method].push error
+  end
+
+  def create_session(request, _unused_call)
+    @requests << request
+    do_create_session(request.database, request.session)
+  end
+
+  def batch_create_sessions(request, _unused_call)
+    @requests << request
+    num_created = 0
+    response = Google::Cloud::Spanner::V1::BatchCreateSessionsResponse.new
+    while num_created < request.session_count
+      response.session << do_create_session(request.database, request.session_template)
+      num_created += 1
+    end
+    response
+  end
+
+  def get_session(request, _unused_call)
+    @requests << request
+    @sessions[request.name]
+  end
+
+  def list_sessions(request, _unused_call)
+    @requests << request
+    response = Google::Cloud::Spanner::V1::ListSessionsResponse.new
+    @sessions.each_value do |s|
+      response.sessions << s
+    end
+    response
+  end
+
+  def delete_session(request, _unused_call)
+    @requests << request
+    @sessions.delete request.name
+    Google::Protobuf::Empty.new
+  end
+
+  def execute_sql(request, _unused_call)
+    do_execute_sql request, false
+  end
+
+  def execute_streaming_sql(request, _unused_call)
+    do_execute_sql request, true
+  end
+
+  # @private
+  def do_execute_sql(request, streaming)
+    @requests << request
+    validate_session request.session
+    created_transaction = do_create_transaction request.session if request.transaction&.begin
+    transaction_id = created_transaction&.id || request.transaction&.id
+    validate_transaction request.session, transaction_id if transaction_id && transaction_id != ""
+
+    raise @errors[request.sql].pop if @errors[request.sql] && !@errors[request.sql].empty?
+
+    result = get_statement_result(request.sql).clone
+    raise result.result if result.result_type == StatementResult::EXCEPTION
+
+    if streaming
+      result.each created_transaction
+    else
+      result.result created_transaction
+    end
+  end
+
+  def execute_batch_dml(request, _unused_call)
+    @requests << request
+    validate_session request.session
+    created_transaction = do_create_transaction request.session if request.transaction&.begin
+    transaction_id = created_transaction&.id || request.transaction&.id
+    validate_transaction request.session, transaction_id if transaction_id && transaction_id != ""
+
+    status = Google::Rpc::Status.new
+    response = Google::Cloud::Spanner::V1::ExecuteBatchDmlResponse.new
+    first = true
+    request.statements.each do |stmt|
+      result = get_statement_result(stmt.sql).clone
+      if result.result_type == StatementResult::EXCEPTION
+        status.code = result.result.code
+        status.message = result.result.message
+        break
+      end
+      if first
+        response.result_sets << result.result(created_transaction)
+        first = false
+      else
+        response.result_sets << result.result
+      end
+    end
+    response.status = status
+    response
+  end
+
+  def read(request, _unused_call)
+    @requests << request
+    raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented"
+  end
+
+  def streaming_read(request, _unused_call)
+    @requests << request
+    raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented"
+  end
+
+  def begin_transaction(request, _unused_call)
+    @requests << request
+    raise @errors[__method__.to_s].pop if @errors[__method__.to_s] && !@errors[__method__.to_s].empty?
+
+    validate_session request.session
+    do_create_transaction request.session
+  end
+
+  def commit(request, _unused_call)
+    @requests << request
+    validate_session request.session
+    validate_transaction request.session, request.transaction_id
+    Google::Cloud::Spanner::V1::CommitResponse.new commit_timestamp: Google::Protobuf::Timestamp.new(seconds: Time.now.to_i)
+  end
+
+  def rollback(request, _unused_call)
+    @requests << request
+    validate_session request.session
+    name = "#{request.session}/transactions/#{request.transaction_id}"
+    @transactions.delete name
+    Google::Protobuf::Empty.new
+  end
+
+  def partition_query(request, _unused_call)
+    @requests << request
+    raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented"
+  end
+
+  def partition_read(request, _unused_call)
+    @requests << request
+    raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented"
+  end
+
+  def get_database(request, _unused_call)
+    @requests << request
+    raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented"
+  end
+
+  def abort_transaction(session, id)
+    return if session.nil? || id.nil?
+
+    name = "#{session}/transactions/#{id}"
+    @aborted_transactions[name] = true
+  end
+
+  def abort_next_transaction
+    @abort_next_transaction = true
+  end
+
+  def get_statement_result(sql)
+    unless @statement_results.key? sql
+      @statement_results.each do |key, value|
+        return value if key.end_with?("%") && sql.start_with?(key.chop)
+      end
+      raise GRPC::BadStatus.new(
+        GRPC::Core::StatusCodes::INVALID_ARGUMENT,
+        "There's no result registered for #{sql}"
+      )
+    end
+    @statement_results[sql]
+  end
+
+  def delete_all_sessions
+    @sessions.clear
+  end
+
+  private
+
+  def validate_session(session)
+    return if @sessions.key? session
+
+    resource_info = Google::Rpc::ResourceInfo.new(
+      resource_type: "type.googleapis.com/google.spanner.v1.Session",
+      resource_name: session
+    )
+    raise GRPC::BadStatus.new(
+      GRPC::Core::StatusCodes::NOT_FOUND,
+      "Session not found: Session with id #{session} not found",
+      { "google.rpc.resourceinfo-bin": Google::Rpc::ResourceInfo.encode(resource_info) }
+    )
+  end
+
+  def do_create_session(database, session_template = nil)
+    @session_counter += 1
+    name = "#{database}/sessions/mock-session-#{@session_counter}"
+    session = Google::Cloud::Spanner::V1::Session.new(name: name)
+
+    session.multiplexed = session_template.multiplexed if session_template
+
+    @sessions[name] = session
+    session
+  end
+
+  def validate_transaction(session, transaction)
+    name = "#{session}/transactions/#{transaction}"
+    unless @transactions.key? name
+      raise GRPC::BadStatus.new(
+        GRPC::Core::StatusCodes::NOT_FOUND,
+        "Transaction not found: Transaction with id #{transaction} not found"
+      )
+    end
+    return unless @aborted_transactions.key?(name)
+
+    retry_info = Google::Rpc::RetryInfo.new(retry_delay: Google::Protobuf::Duration.new(seconds: 0, nanos: 1))
+    raise GRPC::BadStatus.new(
+      GRPC::Core::StatusCodes::ABORTED,
+      "Transaction aborted",
+      { "google.rpc.retryinfo-bin": Google::Rpc::RetryInfo.encode(retry_info) }
+    )
+  end
+
+  def do_create_transaction(session)
+    id = SecureRandom.uuid
+    name = "#{session}/transactions/#{id}"
+    transaction = Google::Cloud::Spanner::V1::Transaction.new id: id
+    @transactions[name] = transaction
+    if @abort_next_transaction
+      abort_transaction session, id
+      @abort_next_transaction = false
+    end
+    transaction
+  end
+end