1 files changed, 52 insertions, 7 deletions

diff --git a/spec/lib/gitlab/spamcheck/client_spec.rb b/spec/lib/gitlab/spamcheck/client_spec.rb
index 956ed2a976f..2fe978125c4 100644
--- a/spec/lib/gitlab/spamcheck/client_spec.rb
+++ b/spec/lib/gitlab/spamcheck/client_spec.rb
@@ -17,6 +17,7 @@ RSpec.describe Gitlab::Spamcheck::Client do
   end
 
   let_it_be(:issue) { create(:issue, description: 'Test issue description') }
+  let_it_be(:snippet) { create(:personal_snippet, :public, description: 'Test issue description') }
 
   let(:response) do
     verdict = ::Spamcheck::SpamVerdict.new
@@ -26,7 +27,7 @@ RSpec.describe Gitlab::Spamcheck::Client do
     verdict
   end
 
-  subject { described_class.new.issue_spam?(spam_issue: issue, user: user) }
+  subject { described_class.new.spam?(spammable: issue, user: user) }
 
   before do
     stub_application_setting(spam_check_endpoint_url: endpoint)
@@ -56,10 +57,11 @@ RSpec.describe Gitlab::Spamcheck::Client do
     end
   end
 
-  describe '#issue_spam?' do
+  shared_examples 'check for spam' do
     before do
       allow_next_instance_of(::Spamcheck::SpamcheckService::Stub) do |instance|
         allow(instance).to receive(:check_for_spam_issue).and_return(response)
+        allow(instance).to receive(:check_for_spam_snippet).and_return(response)
       end
     end
 
@@ -89,12 +91,26 @@ RSpec.describe Gitlab::Spamcheck::Client do
     end
   end
 
-  describe "#build_issue_protobuf", :aggregate_failures do
-    it 'builds the expected protobuf object' do
+  describe "#spam?", :aggregate_failures do
+    describe 'issue' do
+      subject { described_class.new.spam?(spammable: issue, user: user) }
+
+      it_behaves_like "check for spam"
+    end
+
+    describe 'snippet' do
+      subject { described_class.new.spam?(spammable: snippet, user: user, extra_features: { files: [{ path: "file.rb" }] }) }
+
+      it_behaves_like "check for spam"
+    end
+  end
+
+  describe "#build_protobuf", :aggregate_failures do
+    it 'builds the expected issue protobuf object' do
       cxt = { action: :create }
-      issue_pb = described_class.new.send(:build_issue_protobuf,
-                                          issue: issue, user: user,
-                                          context: cxt)
+      issue_pb, _ = described_class.new.send(:build_protobuf,
+                                          spammable: issue, user: user,
+                                          context: cxt, extra_features: {})
       expect(issue_pb.title).to eq issue.title
       expect(issue_pb.description).to eq issue.description
       expect(issue_pb.user_in_project).to be false
@@ -104,6 +120,22 @@ RSpec.describe Gitlab::Spamcheck::Client do
       expect(issue_pb.action).to be ::Spamcheck::Action.lookup(::Spamcheck::Action::CREATE)
       expect(issue_pb.user.username).to eq user.username
     end
+
+    it 'builds the expected snippet protobuf object' do
+      cxt = { action: :create }
+      snippet_pb, _ = described_class.new.send(:build_protobuf,
+                                          spammable: snippet, user: user,
+                                          context: cxt, extra_features: { files: [{ path: 'first.rb' }, { path: 'second.rb' }] })
+      expect(snippet_pb.title).to eq snippet.title
+      expect(snippet_pb.description).to eq snippet.description
+      expect(snippet_pb.created_at).to eq timestamp_to_protobuf_timestamp(snippet.created_at)
+      expect(snippet_pb.updated_at).to eq timestamp_to_protobuf_timestamp(snippet.updated_at)
+      expect(snippet_pb.action).to be ::Spamcheck::Action.lookup(::Spamcheck::Action::CREATE)
+      expect(snippet_pb.user.username).to eq user.username
+      expect(snippet_pb.user.username).to eq user.username
+      expect(snippet_pb.files.first.path).to eq 'first.rb'
+      expect(snippet_pb.files.last.path).to eq 'second.rb'
+    end
   end
 
   describe '#build_user_protobuf', :aggregate_failures do
@@ -143,6 +175,19 @@ RSpec.describe Gitlab::Spamcheck::Client do
     end
   end
 
+  describe "#get_spammable_mappings", :aggregate_failures do
+    it 'is an expected spammable' do
+      protobuf_class, _ = described_class.new.send(:get_spammable_mappings, issue)
+      expect(protobuf_class).to eq ::Spamcheck::Issue
+    end
+
+    it 'is an unexpected spammable' do
+      expect { described_class.new.send(:get_spammable_mappings, 'spam') }.to raise_error(
+        ArgumentError, 'Not a spammable type: String'
+      )
+    end
+  end
+
   private
 
   def timestamp_to_protobuf_timestamp(timestamp)