Address reviewer comments

* create separate method on import_state model

app/models/project_import_state.rb

@@ -6,6 +6,8 @@ class ProjectImportState < ApplicationRecord
 
   self.table_name = "project_mirror_data"
 
+  after_commit :expire_etag_cache
+
   belongs_to :project, inverse_of: :import_state
 
   validates :project, presence: true
@@ -38,16 +40,6 @@ class ProjectImportState < ApplicationRecord
     state :finished
     state :failed
 
-    after_transition any => any do |state|
-      realtime_changes_path = Gitlab::Routing.url_helpers.polymorphic_path([:realtime_changes_import, state.project.import_type.to_sym], format: :json) rescue nil
-
-      if realtime_changes_path
-        Gitlab::EtagCaching::Store.new.tap do |store|
-          store.touch(realtime_changes_path)
-        end
-      end
-    end
-
     after_transition [:none, :finished, :failed] => :scheduled do |state, _|
       state.run_after_commit do
         job_id = project.add_import_job
@@ -88,6 +80,20 @@ class ProjectImportState < ApplicationRecord
     end
   end
 
+  def expire_etag_cache
+    if realtime_changes_path
+      Gitlab::EtagCaching::Store.new.tap do |store|
+        store.touch(realtime_changes_path)
+      end
+    end
+  end
+
+  def realtime_changes_path
+    Gitlab::Routing.url_helpers.polymorphic_path([:realtime_changes_import, project.import_type.to_sym], format: :json)
+  rescue NoMethodError
+    nil
+  end
+
   def relation_hard_failures(limit:)
     project.import_failures.hard_failures_by_correlation_id(correlation_id).limit(limit)
   end

ee/spec/features/registrations/combined_registration_spec.rb

@@ -8,8 +8,8 @@ RSpec.describe 'Combined registration flow', :js do
   let(:experiments) { {} }
 
   before do
-    # https://gitlab.com/gitlab-org/gitlab/-/issues/338737
-    stub_const('Gitlab::QueryLimiting::Transaction::THRESHOLD', 250)
+    # https://gitlab.com/gitlab-org/gitlab/-/issues/340302
+    stub_const('Gitlab::QueryLimiting::Transaction::THRESHOLD', 270)
     stub_experiments(experiments)
     allow(Gitlab).to receive(:com?).and_return(true)
     sign_in(user)

lib/gitlab/github_import/object_counter.rb

@@ -25,7 +25,7 @@ module Gitlab
           increment_project_counter(project, object_type, operation, integer)
           increment_global_counter(object_type, operation, integer)
 
-          expire_etag_cache(project.import_type)
+          project.import_state&.expire_etag_cache
         end
 
         def summary(project)
@@ -44,16 +44,6 @@ module Gitlab
 
         private
 
-        def expire_etag_cache(import_type)
-          realtime_changes_path = Gitlab::Routing.url_helpers.polymorphic_path([:realtime_changes_import, import_type.to_sym], format: :json) rescue nil
-
-          if realtime_changes_path
-            Gitlab::EtagCaching::Store.new.tap do |store|
-              store.touch(realtime_changes_path)
-            end
-          end
-        end
-
         # Global counters are long lived, in Prometheus,
         # and it's used to report the health of the Github Importer
         # in the Grafana Dashboard

spec/controllers/import/github_controller_spec.rb

@@ -258,7 +258,9 @@ RSpec.describe Import::GithubController do
 
         context 'when user input contains colons and spaces' do
           before do
-            allow(controller).to receive(:client_repos).and_return([])
+            allow_next_instance_of(Gitlab::GithubImport::Client) do |client|
+              allow(client).to receive(:search_repos_by_name).and_return(items: [])
+            end
           end
 
           it 'sanitizes user input' do

spec/lib/gitlab/github_import/object_counter_spec.rb

@@ -3,7 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::GithubImport::ObjectCounter, :clean_gitlab_redis_cache do
-  let_it_be(:project) { create(:project, import_type: 'github') }
+  let_it_be(:project) { create(:project, :import_started, import_type: 'github') }
 
   it 'validates the operation being incremented' do
     expect { described_class.increment(project, :issue, :unknown) }

spec/models/project_import_state_spec.rb

@@ -114,47 +114,35 @@ RSpec.describe ProjectImportState, type: :model do
     end
   end
 
-  describe 'import state transitions' do
-    context 'state transition: any => any' do
-      context 'when project import type has realtime changes endpoint' do
-        before do
-          import_state.project.import_type = 'github'
-        end
-
-        it 'expires revelant etag cache' do
-          described_class.state_machines[:status].states.each do |initial_state|
-            import_state.status = initial_state.name
-
-            import_state.status_transitions.each do |transition|
-              expect_next_instance_of(Gitlab::EtagCaching::Store) do |instance|
-                expect(instance).to receive(:touch).with(Gitlab::Routing.url_helpers.realtime_changes_import_github_path(format: :json))
-              end
-
-              transition.perform
-            end
-          end
-        end
+  describe '#expire_etag_cache' do
+    context 'when project import type has realtime changes endpoint' do
+      before do
+        import_state.project.import_type = 'github'
       end
 
-      context 'when project import type has no realtime changes endpoint' do
-        before do
-          import_state.project.import_type = 'jira'
+      it 'expires revelant etag cache' do
+        expect_next_instance_of(Gitlab::EtagCaching::Store) do |instance|
+          expect(instance).to receive(:touch).with(Gitlab::Routing.url_helpers.realtime_changes_import_github_path(format: :json))
         end
 
-        it 'does not touch etag caches' do
-          described_class.state_machines[:status].states.each do |initial_state|
-            import_state.status = initial_state.name
+        subject.expire_etag_cache
+      end
+    end
 
-            import_state.status_transitions.each do |transition|
-              expect(Gitlab::EtagCaching::Store).not_to receive(:new)
+    context 'when project import type does not have realtime changes endpoint' do
+      before do
+        import_state.project.import_type = 'jira'
+      end
 
-              transition.perform
-            end
-          end
-        end
+      it 'does not touch etag caches' do
+        expect(Gitlab::EtagCaching::Store).not_to receive(:new)
+
+        subject.expire_etag_cache
       end
     end
+  end
 
+  describe 'import state transitions' do
     context 'state transition: [:started] => [:finished]' do
       let(:after_import_service) { spy(:after_import_service) }
       let(:housekeeping_service) { spy(:housekeeping_service) }
@@ -231,4 +219,20 @@ RSpec.describe ProjectImportState, type: :model do
       end
     end
   end
+
+  describe 'callbacks' do
+    context 'after_commit :expire_etag_cache' do
+      before do
+        import_state.project.import_type = 'github'
+      end
+
+      it 'expires etag cache' do
+        expect_next_instance_of(Gitlab::EtagCaching::Store) do |instance|
+          expect(instance).to receive(:touch).with(Gitlab::Routing.url_helpers.realtime_changes_import_github_path(format: :json))
+        end
+
+        subject.save!
+      end
+    end
+  end
 end

spec/workers/gitlab/github_import/import_diff_note_worker_spec.rb

@@ -7,7 +7,7 @@ RSpec.describe Gitlab::GithubImport::ImportDiffNoteWorker do
 
   describe '#import' do
     it 'imports a diff note' do
-      project = double(:project, full_path: 'foo/bar', id: 1)
+      project = double(:project, full_path: 'foo/bar', id: 1, import_state: nil)
       client = double(:client)
       importer = double(:importer)
       hash = {

spec/workers/gitlab/github_import/import_issue_worker_spec.rb

@@ -7,7 +7,7 @@ RSpec.describe Gitlab::GithubImport::ImportIssueWorker do
 
   describe '#import' do
     it 'imports an issue' do
-      project = double(:project, full_path: 'foo/bar', id: 1)
+      project = double(:project, full_path: 'foo/bar', id: 1, import_state: nil)
       client = double(:client)
       importer = double(:importer)
       hash = {

spec/workers/gitlab/github_import/import_note_worker_spec.rb

@@ -7,7 +7,7 @@ RSpec.describe Gitlab::GithubImport::ImportNoteWorker do
 
   describe '#import' do
     it 'imports a note' do
-      project = double(:project, full_path: 'foo/bar', id: 1)
+      project = double(:project, full_path: 'foo/bar', id: 1, import_state: nil)
       client = double(:client)
       importer = double(:importer)
       hash = {

spec/workers/gitlab/github_import/import_pull_request_worker_spec.rb

@@ -7,7 +7,7 @@ RSpec.describe Gitlab::GithubImport::ImportPullRequestWorker do
 
   describe '#import' do
     it 'imports a pull request' do
-      project = double(:project, full_path: 'foo/bar', id: 1)
+      project = double(:project, full_path: 'foo/bar', id: 1, import_state: nil)
       client = double(:client)
       importer = double(:importer)
       hash = {