diff --git a/spec/shared/chat_message_classificator_spec.rb b/spec/shared/chat_message_classificator_spec.rb
index 8d2bf59f..30a82308 100644
--- a/spec/shared/chat_message_classificator_spec.rb
+++ b/spec/shared/chat_message_classificator_spec.rb
@@ -15,7 +15,8 @@ describe DiscourseAi::ChatMessageClassificator do
     it "stores the model classification data" do
       classification.classify!(chat_message)
 
-      result = ClassificationResult.find_by(target: chat_message, classification_type: model.type)
+      result =
+        ClassificationResult.find_by(target_id: chat_message.id, classification_type: model.type)
 
       classification = result.classification.symbolize_keys