Add basic tests

Author: presidentbeef

Gemfile

@@ -1,3 +1,7 @@
 source 'https://rubygems.org'
 
 gemspec
+
+group :test do
+  gem 'minitest'
+end

Rakefile

@@ -0,0 +1,9 @@
+require "rake/testtask"
+
+Rake::TestTask.new(:test) do |t|
+  t.libs << "test"
+  t.libs << "lib"
+  t.test_files = FileList["test/test_*.rb"]
+end
+
+task :default => :test

lib/brakeman-llm.rb

@@ -22,6 +22,8 @@ module Brakeman
 
   # Simple wrapper for RubyLLM
   class LLM
+    attr_accessor :instructions, :model, :prompt, :provider
+    attr_reader :llm
 
     # Configure RubyLLM
     def initialize(model:, provider:, instructions: nil, prompt: nil, **kwargs)
@@ -146,18 +148,20 @@ def run(options)
           disclaimer = false
         end
 
-        llm_opts = options.delete(:llm)
+        llm_opts = options.delete(:llm) || {}
 
         # Suppress report output until after analysis
         output_formats = get_output_formats(options)
         output_files = options.delete(:output_files)
         options.delete(:output_format)
         print_report = options.delete(:print_report)
 
+        # Actually run scan
         tracker = old_run(options)
 
+        # Set up LLM
         llm_opts[:log_level] = :debug if @debug
-        llm = Brakeman::LLM.new(**llm_opts)
+        llm = llm_opts.delete(:llm) || Brakeman::LLM.new(**llm_opts)
 
         set_analysis = output_formats.include? :to_json
 
@@ -166,6 +170,7 @@ def run(options)
         warnings = tracker.warnings
         total = warnings.length
 
+        # Update warnings with LLM analysis
         warnings.each_with_index do |warning, index|
           unless @quiet or options[:report_progress] == false
             $stderr.print " #{index}/#{total} warnings processed\r"

test/fixtures/app/app/controllers/application_controller.rb

@@ -0,0 +1,5 @@
+class ApplicationController < ActionController::Base
+  def vulnerable_method
+    eval(params[:x])
+  end
+end

test/test_basics.rb

@@ -0,0 +1,54 @@
+require 'minitest/autorun'
+require 'brakeman-llm'
+require 'brakeman/commandline'
+
+
+class BrakemanLLMTest < Minitest::Test
+  def test_the_basics
+    bm_llm = Brakeman::LLM.new(model: 'test_model', provider: 'test_provider')
+
+    assert_equal 'test_model', bm_llm.model
+    assert_equal 'test_provider', bm_llm.provider
+    assert bm_llm.llm
+
+    # Defaults
+    assert_instance_of String, bm_llm.instructions
+    assert_instance_of String, bm_llm.prompt
+  end
+
+  def test_missing_options
+
+  end
+
+  def test_end_to_end
+    bm_llm = Brakeman::LLM.new(model: 'test_model', provider: 'test_provider')
+    instructions = bm_llm.instructions
+    analysis = "Extended warning description"
+    tracker = nil
+
+    response_mock = Minitest::Mock.new
+    response_mock.expect(:content, analysis.dup)
+    response_mock.expect(:content, analysis.dup)
+
+    chat_mock = Minitest::Mock.new
+
+    2.times do
+      chat_mock.expect(:with_instructions, nil, [instructions])
+      chat_mock.expect(:ask, response_mock, [String])
+    end
+
+    bm_llm.llm.stub(:chat, chat_mock) do
+      tracker = Brakeman.run(llm: { llm: bm_llm }, app_path: File.join(__dir__, 'fixtures', 'app'), output_format: :json)
+    end
+
+    # Check that the warnings have the analysis attached
+    tracker.warnings.each do |w|
+      assert_includes w.llm_analysis, analysis
+
+      assert_equal w.llm_analysis, w.to_hash[:llm_analysis]
+    end
+
+    chat_mock.verify
+    response_mock.verify
+  end
+end