From 73049e982a77610e2573f99ee38920b4e0bfcbec Mon Sep 17 00:00:00 2001
From: Jonathan Claudius <jclaudius@mozilla.com>
Date: Wed, 14 Jun 2017 15:38:46 -0400
Subject: [PATCH] Fix grader bug due to result setting madness

---
 lib/ssh_scan/result.rb       |  4 ++--
 lib/ssh_scan/scan_engine.rb  |  3 +--
 spec/ssh_scan/result_spec.rb | 21 +++++++++++++++++++++
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/lib/ssh_scan/result.rb b/lib/ssh_scan/result.rb
index 2da79939..0e3c0621 100644
--- a/lib/ssh_scan/result.rb
+++ b/lib/ssh_scan/result.rb
@@ -220,11 +220,11 @@ def error
     end
 
     def grade=(grade)
-      @compliance_grade = grade
+      @compliance[:grade] = grade
     end
 
     def grade
-      @compliance_grade
+      @compliance[:grade] 
     end
 
     def to_hash
diff --git a/lib/ssh_scan/scan_engine.rb b/lib/ssh_scan/scan_engine.rb
index 90b167a8..15fb04e4 100644
--- a/lib/ssh_scan/scan_engine.rb
+++ b/lib/ssh_scan/scan_engine.rb
@@ -232,8 +232,7 @@ def scan(opts)
           result.set_compliance = policy_mgr.compliance_results
 
           if result.compliance_policy
-            grader = SSHScan::Grader.new(result)
-            result.grade = grader.grade
+            result.grade = SSHScan::Grader.new(result).grade
           end
         end
       end
diff --git a/spec/ssh_scan/result_spec.rb b/spec/ssh_scan/result_spec.rb
index 9ac9b74f..aa46195c 100644
--- a/spec/ssh_scan/result_spec.rb
+++ b/spec/ssh_scan/result_spec.rb
@@ -191,4 +191,25 @@
       expect(result.compliance_references).to eql(compliance[:references])
     end
   end
+
+  context "when setting grade" do 
+    it "should allow setting of the grade information" do
+      compliance = {
+        :policy => "Test Policy",
+        :compliant => true,
+        :recommendations => ["do this", "do that"],
+        :references => ["https://reference.example.com"],
+      }
+      result = SSHScan::Result.new()
+      result.set_compliance = compliance
+      result.grade = "D"
+
+      expect(result.compliance_policy).to eql(compliance[:policy])
+      expect(result.compliant?).to eql(compliance[:compliant])
+      expect(result.compliance_recommendations).to eql(compliance[:recommendations])
+      expect(result.compliance_references).to eql(compliance[:references])
+      expect(result.grade).to eql("D")
+    end
+  end
+
 end
\ No newline at end of file