Skip to content

Commit

Permalink
exporter_test: cleanup and add more tests
Browse files Browse the repository at this point in the history
  • Loading branch information
zadjadr committed Jul 19, 2024
1 parent 24a27f4 commit 61a8fd5
Show file tree
Hide file tree
Showing 2 changed files with 224 additions and 27 deletions.
250 changes: 223 additions & 27 deletions internal/exporter/exporter_test.go
Original file line number Diff line number Diff line change
@@ -1,10 +1,15 @@
package exporter

import (
"compress/gzip"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"reflect"
"testing"

"github.com/prometheus/client_golang/prometheus"
"zops.top/prometheus-cve-exporter/internal/metrics"
"zops.top/prometheus-cve-exporter/internal/models"
)
Expand Down Expand Up @@ -56,12 +61,127 @@ func TestParsePackagesOutput(t *testing.T) {
}

func TestCheckVulnerabilities(t *testing.T) {
packages := map[string]string{
"vulnerable_package": "1.0.0",
"safe_package": "2.0.0",
testCases := []struct {
name string
packages map[string]string
feed *models.NVDFeed
severity []string
expected map[string]bool
totalVuls float64
}{
{
name: "Vulnerable and safe package",
packages: map[string]string{
"vulnerable_package": "1.0.0",
"safe_package": "2.0.0",
},
feed: createTestFeed("1.0.5"),
severity: []string{"HIGH"},
expected: map[string]bool{
"vulnerable_package": true,
"safe_package": false,
},
totalVuls: 1,
},
{
name: "Both packages safe",
packages: map[string]string{
"vulnerable_package": "1.0.6",
"safe_package": "2.0.0",
},
feed: createTestFeed("1.0.5"),
severity: []string{"HIGH"},
expected: map[string]bool{
"vulnerable_package": false,
"safe_package": false,
},
totalVuls: 0,
},
{
name: "Both packages in vulnerable version range",
packages: map[string]string{
"vulnerable_package": "1.0.0",
"safe_package": "1.0.4",
},
feed: createTestFeed("1.0.5"),
severity: []string{"HIGH"},
expected: map[string]bool{
"vulnerable_package": true,
"safe_package": false,
},
totalVuls: 1,
},
{
name: "Ignore lower severity",
packages: map[string]string{
"vulnerable_package": "1.0.0",
},
feed: createTestFeed("1.0.5"),
severity: []string{"CRITICAL"},
expected: map[string]bool{
"vulnerable_package": false,
},
totalVuls: 0,
},
}

feed := &models.NVDFeed{
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
// Reset metrics before each test case
metrics.ResetVulnerablePackagesGauge()
metrics.SetTotalVulnerabilities(0)

checkVulnerabilities(tc.packages, tc.feed, tc.severity)

metricsGathered, err := prometheus.DefaultGatherer.Gather()
if err != nil {
t.Fatalf("Failed to gather metrics: %v", err)
}

vulnerablePackagesFound := make(map[string]bool)
var totalVulnerabilities float64
var lastUpdateTime float64

for _, metric := range metricsGathered {
switch metric.GetName() {
case "nvd_vulnerable_packages":
for _, m := range metric.GetMetric() {
labels := make(map[string]string)
for _, label := range m.GetLabel() {
labels[label.GetName()] = label.GetValue()
}
pkg := labels["package"]
vulnerablePackagesFound[pkg] = true
if m.GetGauge().GetValue() != 1 {
t.Errorf("Expected vulnerable package metric to be 1 for %s, got %v", pkg, m.GetGauge().GetValue())
}
}
case "nvd_total_vulnerabilities":
totalVulnerabilities = metric.GetMetric()[0].GetGauge().GetValue()
case "nvd_last_update_time":
lastUpdateTime = metric.GetMetric()[0].GetGauge().GetValue()
}
}

for pkg, expectedVulnerable := range tc.expected {
if expectedVulnerable != vulnerablePackagesFound[pkg] {
t.Errorf("Package %s: expected vulnerable = %v, got %v", pkg, expectedVulnerable, vulnerablePackagesFound[pkg])
}
}

if totalVulnerabilities != tc.totalVuls {
t.Errorf("Expected total vulnerabilities to be %v, got %v", tc.totalVuls, totalVulnerabilities)
}

if lastUpdateTime <= 0 {
t.Errorf("Expected last update time to be > 0, got %v", lastUpdateTime)
}
})
}
}

func createTestFeed(versionEndExcluding string) *models.NVDFeed {
return &models.NVDFeed{
CVEItems: []models.CVEItem{
{
CVE: models.CVE{
Expand All @@ -81,8 +201,9 @@ func TestCheckVulnerabilities(t *testing.T) {
{
CPEMatch: []models.CPEMatch{
{
Vulnerable: true,
CPE23Uri: "cpe:2.3:a:vendor:vulnerable_package:1.0.0:*:*:*:*:*:*:*",
Vulnerable: true,
CPE23Uri: "cpe:2.3:a:vendor:vulnerable_package:1.0.0:*:*:*:*:*:*:*",
VersionEndExcluding: versionEndExcluding,
},
},
},
Expand All @@ -91,17 +212,6 @@ func TestCheckVulnerabilities(t *testing.T) {
},
},
}

severity := []string{"HIGH"}

// Reset metrics before test
metrics.ResetVulnerablePackagesGauge()

checkVulnerabilities(packages, feed, severity)

// Here you would typically check the metrics.
// For this example, we'll just verify that the function runs without panicking.
// In a real test, you'd use a mocked metrics package to verify the correct metrics were set.
}

func TestIsVersionVulnerable(t *testing.T) {
Expand Down Expand Up @@ -146,17 +256,103 @@ func TestIsVersionVulnerable(t *testing.T) {
}

func TestFetchNVDFeed(t *testing.T) {
// This is a mock test. In a real scenario, you'd use a mock HTTP server.
t.Skip("Skipping TestFetchNVDFeed as it requires network access")

feed, err := fetchNVDFeed("https://example.com/feed.json")
if err != nil {
t.Fatalf("fetchNVDFeed() error = %v", err)
tests := []struct {
name string
cveItems []models.CVEItem
expectedID string
expectedLength int
expectError bool
}{
{
name: "single item",
cveItems: []models.CVEItem{
{
CVE: models.CVE{
CVEDataMeta: models.CVEDataMeta{
ID: "CVE-2023-TEST",
},
},
},
},
expectedID: "CVE-2023-TEST",
expectedLength: 1,
expectError: false,
},
{
name: "multiple items",
cveItems: []models.CVEItem{
{
CVE: models.CVE{
CVEDataMeta: models.CVEDataMeta{
ID: "CVE-2023-TEST-1",
},
},
},
{
CVE: models.CVE{
CVEDataMeta: models.CVEDataMeta{
ID: "CVE-2023-TEST-2",
},
},
},
},
expectedID: "CVE-2023-TEST-1",
expectedLength: 2,
expectError: false,
},
{
name: "empty feed",
cveItems: []models.CVEItem{},
expectedID: "",
expectedLength: 0,
expectError: false,
},
}

if feed == nil {
t.Error("fetchNVDFeed() returned nil feed")
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Create a test server
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Encoding", "gzip")

gzWriter := gzip.NewWriter(w)
defer gzWriter.Close()

// Add more specific checks on the feed content if needed
// Create a JSON encoder and encode the data
encoder := json.NewEncoder(gzWriter)
err := encoder.Encode(models.NVDFeed{
CVEItems: tt.cveItems,
})
if err != nil {
t.Fatalf("Failed to encode JSON: %v", err)
}

// Ensure that the gzip writer is flushed
if err := gzWriter.Close(); err != nil {
t.Fatalf("Failed to close gzip writer: %v", err)
}
}))
defer server.Close()

feed, err := fetchNVDFeed(server.URL)
if (err != nil) != tt.expectError {
t.Fatalf("fetchNVDFeed() error = %v, wantErr %v", err, tt.expectError)
}

if feed == nil {
if !tt.expectError {
t.Error("fetchNVDFeed() returned nil feed")
}
return
}

if feed.CVEItems == nil {
t.Error("feed.CVEItems is nil")
} else if len(feed.CVEItems) != tt.expectedLength {
t.Errorf("expected %d CVEItems, got %d", tt.expectedLength, len(feed.CVEItems))
} else if len(feed.CVEItems) > 0 && feed.CVEItems[0].CVE.CVEDataMeta.ID != tt.expectedID {
t.Errorf("expected CVE ID '%s', got %s", tt.expectedID, feed.CVEItems[0].CVE.CVEDataMeta.ID)
}
})
}
}
1 change: 1 addition & 0 deletions internal/metrics/metrics.go
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ func SetLastUpdateTime() {
lastUpdateTimeGauge.SetToCurrentTime()
}

// ResetVulnerablePackagesGauge resets the vulnerablePackagesGauge
func ResetVulnerablePackagesGauge() {
vulnerablePackagesGauge.Reset()
}
Expand Down

0 comments on commit 61a8fd5

Please sign in to comment.