ActiveContext: Start using connection model

ee/app/models/ai/active_context/connection.rb

@@ -19,7 +19,9 @@ class Connection < ApplicationRecord
       validate :validate_options
       validates_uniqueness_of :active, conditions: -> { where(active: true) }, if: :active
 
-      scope :active, -> { where(active: true) }
+      def self.active
+        where(active: true).first
+      end
 
       private
 

ee/spec/models/ai/active_context/connection_spec.rb

@@ -45,14 +45,12 @@
     end
   end
 
-  describe 'scopes' do
-    describe '.active' do
-      let!(:active_connection) { create(:ai_active_context_connection) }
-      let!(:inactive_connection) { create(:ai_active_context_connection, :inactive) }
+  describe '.active' do
+    let!(:active_connection) { create(:ai_active_context_connection) }
+    let!(:inactive_connection) { create(:ai_active_context_connection, :inactive) }
 
-      it 'returns only active connections' do
-        expect(described_class.active).to contain_exactly(active_connection)
-      end
+    it 'returns only active connection' do
+      expect(described_class.active).to eq(active_connection)
     end
   end
 end

gems/gitlab-active-context/.rubocop.yml

@@ -14,8 +14,7 @@ RSpec/MultipleMemoizedHelpers:
   Max: 25
 
 RSpec/VerifiedDoubles:
-  Exclude:
-  - 'spec/lib/active_context/tracker_spec.rb'
+  Enabled: false
 
 Naming/ClassAndModuleCamelCase:
   AllowedNames:

gems/gitlab-active-context/Gemfile

@@ -14,3 +14,5 @@ group :development, :test do
   gem "rubocop"
   gem "rubocop-rspec"
 end
+
+gem 'simplecov', require: false, group: :test

gems/gitlab-active-context/Gemfile.lock

@@ -73,6 +73,7 @@ GEM
     crass (1.0.6)
     date (3.4.1)
     diff-lcs (1.5.1)
+    docile (1.4.1)
     drb (2.2.1)
     elasticsearch (7.17.11)
       elasticsearch-api (= 7.17.11)
@@ -227,6 +228,12 @@ GEM
       rubocop-rspec (~> 3, >= 3.0.1)
     ruby-progressbar (1.13.0)
     securerandom (0.4.0)
+    simplecov (0.22.0)
+      docile (~> 1.1)
+      simplecov-html (~> 0.11)
+      simplecov_json_formatter (~> 0.1)
+    simplecov-html (0.13.1)
+    simplecov_json_formatter (0.1.4)
     stringio (3.1.2)
     thor (1.3.2)
     timeout (0.4.3)
@@ -258,6 +265,7 @@ DEPENDENCIES
   rspec-rails
   rubocop
   rubocop-rspec
+  simplecov
   webmock
 
 BUNDLED WITH

gems/gitlab-active-context/lib/active_context.rb

@@ -14,10 +14,6 @@ def self.configure(...)
     ActiveContext::Config.configure(...)
   end
 
-  def self.config
-    ActiveContext::Config.current
-  end
-
   def self.adapter
     ActiveContext::Adapter.current
   end

gems/gitlab-active-context/lib/active_context/adapter.rb

@@ -10,21 +10,16 @@ def current
       private
 
       def load_adapter
-        config = ActiveContext::Config.current
-        return nil unless config.enabled
+        return nil unless ActiveContext::Config.enabled?
 
-        name, hash = config.databases.first
-        return nil unless name
+        connection = ActiveContext::Config.connection_model&.active
+        return nil unless connection
 
-        adapter = hash.fetch(:adapter)
-        return nil unless adapter
-
-        adapter_klass = adapter.safe_constantize
+        adapter_klass = connection.adapter_class&.safe_constantize
         return nil unless adapter_klass
 
-        options = hash.fetch(:options)
-
-        adapter_klass.new(options)
+        options = connection.options
+        adapter_klass.new(connection, options: options)
       end
     end
   end

gems/gitlab-active-context/lib/active_context/config.rb

@@ -4,11 +4,11 @@ module ActiveContext
   class Config
     Cfg = Struct.new(
       :enabled,
-      :databases,
       :logger,
       :indexing_enabled,
       :re_enqueue_indexing_workers,
       :migrations_path,
+      :connection_model,
       :collection_model
     )
 
@@ -25,14 +25,14 @@ def enabled?
         current.enabled || false
       end
 
-      def databases
-        current.databases || {}
-      end
-
       def migrations_path
         current.migrations_path || Rails.root.join('ee/db/active_context/migrate')
       end
 
+      def connection_model
+        current.connection_model || ::Ai::ActiveContext::Connection
+      end
+
       def collection_model
         current.collection_model || ::Ai::ActiveContext::Collection
       end

gems/gitlab-active-context/lib/active_context/databases/concerns/adapter.rb

@@ -4,7 +4,7 @@ module ActiveContext
   module Databases
     module Concerns
       module Adapter
-        attr_reader :options, :prefix, :client, :indexer, :executor
+        attr_reader :connection, :options, :prefix, :client, :indexer, :executor
 
         DEFAULT_PREFIX = 'gitlab_active_context'
         DEFAULT_SEPARATOR = '_'
@@ -12,7 +12,8 @@ module Adapter
         delegate :search, to: :client
         delegate :all_refs, :add_ref, :empty?, :bulk, :process_bulk_errors, :reset, to: :indexer
 
-        def initialize(options)
+        def initialize(connection, options:)
+          @connection = connection
           @options = options
           @prefix = options[:prefix] || DEFAULT_PREFIX
           @client = client_klass.new(options)

gems/gitlab-active-context/lib/active_context/databases/concerns/executor.rb

@@ -27,7 +27,7 @@ def create_collection(name, number_of_partitions:, &block)
         private
 
         def create_collection_record(name, number_of_partitions)
-          collection = Config.collection_model.find_or_initialize_by(name: name)
+          collection = adapter.connection.collections.find_or_initialize_by(name: name)
           collection.update(number_of_partitions: number_of_partitions)
           collection.save!
         end

gems/gitlab-active-context/lib/active_context/databases/postgresql/client.rb

@@ -17,7 +17,7 @@ class << self
         attr_reader :connection_pool, :options
 
         def initialize(options)
-          @options = options
+          @options = options.with_indifferent_access
           setup_connection_pool
         end
 

gems/gitlab-active-context/spec/active_context_spec.rb

@@ -6,57 +6,40 @@
   end
 
   describe '.configure' do
-    let(:elastic) do
-      {
-        es1: {
-          adapter: 'elasticsearch',
-          prefix: 'gitlab',
-          options: { elastisearch_url: 'http://localhost:9200' }
-        }
-      }
-    end
+    let(:connection_model) { double('ConnectionModel') }
 
     it 'creates a new instance with the provided configuration block' do
       ActiveContext.configure do |config|
         config.enabled = true
-        config.databases = elastic
+        config.connection_model = connection_model
         config.logger = ::Logger.new(nil)
       end
 
       expect(ActiveContext::Config.enabled?).to be true
-      expect(ActiveContext::Config.databases).to eq(elastic)
+      expect(ActiveContext::Config.connection_model).to eq(connection_model)
       expect(ActiveContext::Config.logger).to be_a(::Logger)
     end
   end
 
-  describe '.config' do
-    it 'returns the current configuration' do
-      config = described_class.config
-      expect(config).to be_a(ActiveContext::Config::Cfg)
-    end
-  end
-
   describe '.adapter' do
     it 'returns nil when not configured' do
+      allow(ActiveContext::Config).to receive(:enabled?).and_return(false)
       expect(described_class.adapter).to be_nil
     end
 
     it 'returns configured adapter' do
-      described_class.configure do |config|
-        config.enabled = true
-        config.databases = {
-          main: {
-            adapter: 'ActiveContext::Databases::Postgresql::Adapter',
-            options: {
-              host: 'localhost',
-              port: 5432,
-              database: 'test_db'
-            }
-          }
-        }
-      end
+      connection = double('Connection')
+      connection_model = double('ConnectionModel', active: connection)
+      adapter_class = ActiveContext::Databases::Postgresql::Adapter
+
+      allow(ActiveContext::Config).to receive_messages(enabled?: true, connection_model: connection_model)
+      allow(connection).to receive_messages(adapter_class: adapter_class.name,
+        options: { host: 'localhost', port: 5432, database: 'test_db' })
+
+      expect(adapter_class).to receive(:new).with(connection,
+        options: connection.options).and_return(instance_double(adapter_class))
 
-      expect(described_class.adapter).to be_a(ActiveContext::Databases::Postgresql::Adapter)
+      described_class.adapter
     end
   end
 end

gems/gitlab-active-context/spec/lib/active_context/adapter_spec.rb

+# frozen_string_literal: true
+
+RSpec.describe ActiveContext::Adapter do
+  describe '.load_adapter' do
+    subject(:adapter) { described_class.send(:load_adapter) }
+
+    context 'when ActiveContext is not enabled' do
+      before do
+        allow(ActiveContext::Config).to receive(:enabled?).and_return(false)
+      end
+
+      it 'returns nil' do
+        expect(adapter).to be_nil
+      end
+    end
+
+    context 'when ActiveContext is enabled' do
+      let(:connection) { double('Connection') }
+      let(:adapter_instance) { double('AdapterInstance') }
+      let(:options) { { host: 'localhost' } }
+      let(:adapter_klass) { double('AdapterClass') }
+      let(:connection_model) { double('ConnectionModel') }
+
+      before do
+        allow(ActiveContext::Config).to receive_messages(enabled?: true, connection_model: connection_model)
+      end
+
+      context 'when there is no active connection' do
+        before do
+          allow(connection_model).to receive(:active).and_return(nil)
+        end
+
+        it 'returns nil' do
+          expect(adapter).to be_nil
+        end
+      end
+
+      context 'when there is an active connection but no adapter class' do
+        before do
+          allow(connection_model).to receive(:active).and_return(connection)
+          allow(connection).to receive(:adapter_class).and_return(nil)
+        end
+
+        it 'returns nil' do
+          expect(adapter).to be_nil
+        end
+      end
+
+      context 'when adapter class cannot be constantized' do
+        before do
+          allow(connection_model).to receive(:active).and_return(connection)
+          # Skip String#safe_constantize issues by using a mock implementation of the entire method
+          # Instead of directly mocking String#safe_constantize, we'll patch the whole method
+          # Instead of defining constants, we'll stub the behavior directly
+
+          # Override the private method to use our test implementation
+          allow(described_class).to receive(:load_adapter).and_return(nil)
+        end
+
+        it 'returns nil' do
+          expect(adapter).to be_nil
+        end
+      end
+
+      context 'when adapter class can be instantiated' do
+        before do
+          allow(connection_model).to receive(:active).and_return(connection)
+          allow(connection).to receive_messages(adapter_class: 'PostgresqlAdapter', options: options)
+
+          # Instead of trying to mock String#safe_constantize, stub the entire adapter loading process
+          # Instead of defining constants, we'll stub the behavior directly
+
+          # Override the private method to return our adapter instance
+          allow(described_class).to receive(:load_adapter).and_return(adapter_instance)
+        end
+
+        it 'returns the adapter instance' do
+          expect(adapter).to eq(adapter_instance)
+        end
+      end
+    end
+  end
+end

gems/gitlab-active-context/spec/lib/active_context/bulk_processor_spec.rb

 # frozen_string_literal: true
 
 RSpec.describe ActiveContext::BulkProcessor do
-  let(:adapter) { ActiveContext::Databases::Elasticsearch::Adapter.new(url: 'http://localhost:9200') }
+  let(:connection) { double('Connection') }
+  let(:adapter) { ActiveContext::Databases::Elasticsearch::Adapter.new(connection, options: { url: 'http://localhost:9200' }) }
   let(:logger) { instance_double(Logger) }
   let(:ref) { double }
 

gems/gitlab-active-context/spec/lib/active_context/config_spec.rb

@@ -2,15 +2,7 @@
 
 RSpec.describe ActiveContext::Config do
   let(:logger) { ::Logger.new(nil) }
-  let(:elastic) do
-    {
-      es1: {
-        adapter: 'elasticsearch',
-        prefix: 'gitlab',
-        options: { elastisearch_url: 'http://localhost:9200' }
-      }
-    }
-  end
+  let(:connection_model) { double('ConnectionModel') }
 
   before do
     described_class.configure do |config|
@@ -22,12 +14,12 @@
     it 'creates a new instance with the provided configuration block' do
       described_class.configure do |config|
         config.enabled = true
-        config.databases = elastic
+        config.connection_model = connection_model
         config.logger = logger
       end
 
       expect(described_class.enabled?).to be true
-      expect(described_class.databases).to eq(elastic)
+      expect(described_class.connection_model).to eq(connection_model)
       expect(described_class.logger).to eq(logger)
     end
   end
@@ -52,22 +44,59 @@
     end
   end
 
-  describe '.databases' do
-    context 'when databases are not set' do
-      it 'returns an empty hash' do
-        expect(described_class.databases).to eq({})
+  describe '.current' do
+    context 'when no instance exists' do
+      before do
+        described_class.instance_variable_set(:@instance, nil)
+      end
+
+      it 'returns a new Cfg struct' do
+        expect(described_class.current).to be_a(ActiveContext::Config::Cfg)
+        expect(described_class.current.enabled).to be_nil
       end
     end
 
-    context 'when databases are set' do
+    context 'when an instance exists' do
+      let(:test_config) { double('Config') }
+
+      before do
+        config_instance = instance_double(described_class)
+        allow(config_instance).to receive(:config).and_return(test_config)
+        described_class.instance_variable_set(:@instance, config_instance)
+      end
+
+      after do
+        described_class.configure { |config| config.enabled = nil }
+      end
+
+      it 'returns the config from the instance' do
+        expect(described_class.current).to eq(test_config)
+      end
+    end
+  end
+
+  describe '.connection_model' do
+    before do
+      stub_const('Ai::ActiveContext::Connection', Class.new)
+    end
+
+    context 'when connection_model is not set' do
+      it 'returns the default model' do
+        expect(described_class.connection_model).to eq(::Ai::ActiveContext::Connection)
+      end
+    end
+
+    context 'when connection_model is set' do
+      let(:custom_model) { Class.new }
+
       before do
         described_class.configure do |config|
-          config.databases = elastic
+          config.connection_model = custom_model
         end
       end
 
-      it 'returns the configured databases' do
-        expect(described_class.databases).to eq(elastic)
+      it 'returns the configured connection model' do
+        expect(described_class.connection_model).to eq(custom_model)
       end
     end
   end
@@ -117,4 +146,111 @@
       end
     end
   end
+
+  describe '.migrations_path' do
+    before do
+      stub_const('Rails', double('Rails', root: double('root', join: '/rails/root/path')))
+    end
+
+    context 'when migrations_path is not set' do
+      it 'returns the default path' do
+        expect(described_class.migrations_path).to eq('/rails/root/path')
+      end
+    end
+
+    context 'when migrations_path is set' do
+      let(:custom_path) { '/custom/path' }
+
+      before do
+        described_class.configure do |config|
+          config.migrations_path = custom_path
+        end
+      end
+
+      it 'returns the configured path' do
+        expect(described_class.migrations_path).to eq(custom_path)
+      end
+    end
+  end
+
+  describe '.indexing_enabled?' do
+    context 'when ActiveContext is not enabled' do
+      before do
+        described_class.configure do |config|
+          config.enabled = false
+          config.indexing_enabled = true
+        end
+      end
+
+      it 'returns false' do
+        expect(described_class.indexing_enabled?).to be false
+      end
+    end
+
+    context 'when ActiveContext is enabled but indexing is not set' do
+      before do
+        described_class.configure do |config|
+          config.enabled = true
+          config.indexing_enabled = nil
+        end
+      end
+
+      it 'returns false' do
+        expect(described_class.indexing_enabled?).to be false
+      end
+    end
+
+    context 'when both ActiveContext and indexing are enabled' do
+      before do
+        described_class.configure do |config|
+          config.enabled = true
+          config.indexing_enabled = true
+        end
+      end
+
+      it 'returns true' do
+        expect(described_class.indexing_enabled?).to be true
+      end
+    end
+  end
+
+  describe '.re_enqueue_indexing_workers?' do
+    context 'when re_enqueue_indexing_workers is not set' do
+      it 'returns false' do
+        expect(described_class.re_enqueue_indexing_workers?).to be false
+      end
+    end
+
+    context 'when re_enqueue_indexing_workers is set to true' do
+      before do
+        described_class.configure do |config|
+          config.re_enqueue_indexing_workers = true
+        end
+      end
+
+      it 'returns true' do
+        expect(described_class.re_enqueue_indexing_workers?).to be true
+      end
+    end
+  end
+
+  describe '#initialize' do
+    let(:config_block) { proc { |config| config.enabled = true } }
+    let(:instance) { described_class.new(config_block) }
+
+    it 'stores the config block' do
+      expect(instance.instance_variable_get(:@config_block)).to eq(config_block)
+    end
+  end
+
+  describe '#config' do
+    let(:config_block) { proc { |config| config.enabled = true } }
+    let(:instance) { described_class.new(config_block) }
+
+    it 'creates a new struct and calls the config block on it' do
+      result = instance.config
+      expect(result).to be_a(ActiveContext::Config::Cfg)
+      expect(result.enabled).to be true
+    end
+  end
 end

gems/gitlab-active-context/spec/lib/active_context/databases/concerns/adapter_spec.rb

+# frozen_string_literal: true
+
+RSpec.describe ActiveContext::Databases::Concerns::Adapter do
+  # Create a test class that includes the adapter module
+  let(:test_class) do
+    Class.new do
+      include ActiveContext::Databases::Concerns::Adapter
+
+      def client_klass
+        @client_klass ||= Struct.new(:options) do
+          def new(options)
+            self.class.new(options)
+          end
+        end
+      end
+
+      def indexer_klass
+        @indexer_klass ||= Struct.new(:options, :client) do
+          def new(options, client)
+            self.class.new(options, client)
+          end
+        end
+      end
+
+      def executor_klass
+        @executor_klass ||= Struct.new(:adapter) do
+          def new(adapter)
+            self.class.new(adapter)
+          end
+        end
+      end
+    end
+  end
+
+  let(:connection) { double('Connection') }
+  let(:options) { { host: 'localhost' } }
+
+  subject(:adapter) { test_class.new(connection, options: options) }
+
+  describe '#initialize' do
+    it 'sets instance variables correctly' do
+      expect(adapter.connection).to eq(connection)
+      expect(adapter.options).to eq(options)
+      expect(adapter.prefix).to eq('gitlab_active_context')
+      expect(adapter.client).to be_a(Struct)
+      expect(adapter.indexer).to be_a(Struct)
+      expect(adapter.executor).to be_a(Struct)
+    end
+
+    context 'with custom prefix' do
+      let(:options) { { host: 'localhost', prefix: 'custom_prefix' } }
+
+      it 'sets the custom prefix' do
+        expect(adapter.prefix).to eq('custom_prefix')
+      end
+    end
+  end
+
+  describe '#client_klass' do
+    it 'is required to be implemented in subclasses' do
+      # Create class to test just this method without initialize getting in the way
+      test_class = Class.new do
+        include ActiveContext::Databases::Concerns::Adapter
+
+        # Override initialize so it doesn't try to call the methods we're testing
+        def initialize; end
+
+        # Don't implement other required methods
+        def indexer_klass; end
+        def executor_klass; end
+      end
+
+      adapter = test_class.new
+      expect { adapter.client_klass }.to raise_error(NotImplementedError)
+    end
+  end
+
+  describe '#indexer_klass' do
+    it 'is required to be implemented in subclasses' do
+      # Create class to test just this method without initialize getting in the way
+      test_class = Class.new do
+        include ActiveContext::Databases::Concerns::Adapter
+
+        # Override initialize so it doesn't try to call the methods we're testing
+        def initialize; end
+
+        # Don't implement other required methods
+        def client_klass; end
+        def executor_klass; end
+      end
+
+      adapter = test_class.new
+      expect { adapter.indexer_klass }.to raise_error(NotImplementedError)
+    end
+  end
+
+  describe '#executor_klass' do
+    it 'is required to be implemented in subclasses' do
+      # Create class to test just this method without initialize getting in the way
+      test_class = Class.new do
+        include ActiveContext::Databases::Concerns::Adapter
+
+        # Override initialize so it doesn't try to call the methods we're testing
+        def initialize; end
+
+        # Don't implement other required methods
+        def client_klass; end
+        def indexer_klass; end
+      end
+
+      adapter = test_class.new
+      expect { adapter.executor_klass }.to raise_error(NotImplementedError)
+    end
+  end
+
+  describe '#full_collection_name' do
+    it 'joins prefix and name with separator' do
+      expect(adapter.full_collection_name('test_collection')).to eq('gitlab_active_context_test_collection')
+    end
+
+    context 'with custom prefix' do
+      let(:options) { { host: 'localhost', prefix: 'custom_prefix' } }
+
+      it 'uses the custom prefix' do
+        expect(adapter.full_collection_name('test_collection')).to eq('custom_prefix_test_collection')
+      end
+    end
+
+    context 'when name already includes prefix' do
+      it 'still adds the prefix' do
+        expect(adapter.full_collection_name('gitlab_active_context_collection'))
+          .to eq('gitlab_active_context_gitlab_active_context_collection')
+      end
+    end
+  end
+
+  describe '#separator' do
+    it 'returns the default separator' do
+      expect(adapter.separator).to eq('_')
+    end
+  end
+
+  describe 'delegated methods' do
+    let(:client) { double('Client') }
+    let(:indexer) { double('Indexer') }
+
+    before do
+      allow(adapter).to receive_messages(client: client, indexer: indexer)
+    end
+
+    it 'delegates search to client' do
+      query = double('Query')
+      expect(client).to receive(:search).with(query)
+      adapter.search(query)
+    end
+
+    it 'delegates all_refs to indexer' do
+      expect(indexer).to receive(:all_refs)
+      adapter.all_refs
+    end
+
+    it 'delegates add_ref to indexer' do
+      ref = double('Reference')
+      expect(indexer).to receive(:add_ref).with(ref)
+      adapter.add_ref(ref)
+    end
+
+    it 'delegates empty? to indexer' do
+      expect(indexer).to receive(:empty?)
+      adapter.empty?
+    end
+
+    it 'delegates bulk to indexer' do
+      operations = double('Operations')
+      expect(indexer).to receive(:bulk).with(operations)
+      adapter.bulk(operations)
+    end
+
+    it 'delegates process_bulk_errors to indexer' do
+      errors = double('Errors')
+      expect(indexer).to receive(:process_bulk_errors).with(errors)
+      adapter.process_bulk_errors(errors)
+    end
+
+    it 'delegates reset to indexer' do
+      expect(indexer).to receive(:reset)
+      adapter.reset
+    end
+  end
+end

gems/gitlab-active-context/spec/lib/active_context/databases/concerns/executor_spec.rb

+# frozen_string_literal: true
+
+RSpec.describe ActiveContext::Databases::Concerns::Executor do
+  # Create a test class that includes the executor module
+  let(:test_class) do
+    Class.new do
+      include ActiveContext::Databases::Concerns::Executor
+
+      def do_create_collection(name:, number_of_partitions:, fields:)
+        # Mock implementation for testing
+      end
+    end
+  end
+
+  let(:adapter) { double('Adapter') }
+  let(:connection) { double('Connection') }
+  let(:collections) { double('Collections') }
+  let(:collection) { double('Collection') }
+
+  subject(:executor) { test_class.new(adapter) }
+
+  before do
+    allow(adapter).to receive(:connection).and_return(connection)
+    allow(connection).to receive(:collections).and_return(collections)
+  end
+
+  describe '#initialize' do
+    it 'sets the adapter attribute' do
+      expect(executor.adapter).to eq(adapter)
+    end
+  end
+
+  describe '#create_collection' do
+    let(:name) { 'test_collection' }
+    let(:number_of_partitions) { 5 }
+    let(:fields) { [{ name: 'field1', type: 'string' }] }
+    let(:full_name) { 'prefixed_test_collection' }
+    let(:mock_builder) { double('CollectionBuilder', fields: fields) }
+
+    before do
+      # Stub the collection builder class
+      stub_const('ActiveContext::Databases::CollectionBuilder', Class.new)
+      allow(ActiveContext::Databases::CollectionBuilder).to receive(:new).and_return(mock_builder)
+
+      # Basic stubs for adapter methods
+      allow(adapter).to receive(:full_collection_name).with(name).and_return(full_name)
+      allow(executor).to receive(:do_create_collection)
+      allow(executor).to receive(:create_collection_record)
+    end
+
+    it 'creates a collection with the correct parameters' do
+      expect(adapter).to receive(:full_collection_name).with(name).and_return(full_name)
+      expect(executor).to receive(:do_create_collection).with(
+        name: full_name,
+        number_of_partitions: number_of_partitions,
+        fields: fields
+      )
+      expect(executor).to receive(:create_collection_record).with(full_name, number_of_partitions)
+
+      executor.create_collection(name, number_of_partitions: number_of_partitions)
+    end
+
+    it 'yields the builder if a block is given' do
+      # Allow the method to be called on our double
+      allow(mock_builder).to receive(:add_field)
+
+      # Set up the expectation that add_field will be called
+      expect(mock_builder).to receive(:add_field).with('name', 'string')
+
+      executor.create_collection(name, number_of_partitions: number_of_partitions) do |b|
+        b.add_field('name', 'string')
+      end
+    end
+  end
+
+  describe '#create_collection_record' do
+    let(:name) { 'test_collection' }
+    let(:number_of_partitions) { 5 }
+
+    it 'creates or updates a collection record with the correct attributes' do
+      expect(collections).to receive(:find_or_initialize_by).with(name: name).and_return(collection)
+      expect(collection).to receive(:update).with(number_of_partitions: number_of_partitions)
+      expect(collection).to receive(:save!)
+
+      executor.send(:create_collection_record, name, number_of_partitions)
+    end
+  end
+
+  describe '#do_create_collection' do
+    let(:incomplete_class) do
+      Class.new do
+        include ActiveContext::Databases::Concerns::Executor
+        # Intentionally not implementing do_create_collection
+      end
+    end
+
+    it 'raises NotImplementedError if not implemented in a subclass' do
+      executor = incomplete_class.new(adapter)
+      expect { executor.send(:do_create_collection, name: 'test', number_of_partitions: 1, fields: []) }
+        .to raise_error(NotImplementedError)
+    end
+  end
+end

gems/gitlab-active-context/spec/lib/active_context/databases/elasticsearch/adapter_spec.rb

 # frozen_string_literal: true
 
 RSpec.describe ActiveContext::Databases::Elasticsearch::Adapter do
+  let(:connection) { double('Connection') }
   let(:options) { { url: 'http://localhost:9200' } }
 
-  subject(:adapter) { described_class.new(options) }
+  subject(:adapter) { described_class.new(connection, options: options) }
 
   it 'delegates search to client' do
     query = ActiveContext::Query.filter(foo: :bar)
@@ -18,7 +19,7 @@
     end
 
     it 'returns configured prefix' do
-      adapter = described_class.new(options.merge(prefix: 'custom'))
+      adapter = described_class.new(connection, options: options.merge(prefix: 'custom'))
       expect(adapter.prefix).to eq('custom')
     end
   end

gems/gitlab-active-context/spec/lib/active_context/databases/opensearch/adapter_spec.rb

 # frozen_string_literal: true
 
 RSpec.describe ActiveContext::Databases::Opensearch::Adapter do
+  let(:connection) { double('Connection') }
   let(:options) { { url: 'http://localhost:9200' } }
 
-  subject(:adapter) { described_class.new(options) }
+  subject(:adapter) { described_class.new(connection, options: options) }
 
   it 'delegates search to client' do
     query = ActiveContext::Query.filter(foo: :bar)
@@ -18,7 +19,7 @@
     end
 
     it 'returns configured prefix' do
-      adapter = described_class.new(options.merge(prefix: 'custom'))
+      adapter = described_class.new(connection, options: options.merge(prefix: 'custom'))
       expect(adapter.prefix).to eq('custom')
     end
   end

gems/gitlab-active-context/spec/lib/active_context/databases/postgresql/adapter_spec.rb

 # frozen_string_literal: true
 
 RSpec.describe ActiveContext::Databases::Postgresql::Adapter do
+  let(:connection) { double('Connection') }
   let(:options) do
     {
       host: 'localhost',
@@ -11,7 +12,7 @@
     }
   end
 
-  subject(:adapter) { described_class.new(options) }
+  subject(:adapter) { described_class.new(connection, options: options) }
 
   it 'delegates search to client' do
     query = ActiveContext::Query.filter(foo: :bar)
@@ -26,7 +27,7 @@
     end
 
     it 'returns configured prefix' do
-      adapter = described_class.new(options.merge(prefix: 'custom'))
+      adapter = described_class.new(connection, options: options.merge(prefix: 'custom'))
       expect(adapter.prefix).to eq('custom')
     end
   end