Modify NewNVDFetcher to allow setting NVDKey separately

mtnmunuklu · Feb 13, 2024 · bb7a4a5 · bb7a4a5
1 parent a7a1006
commit bb7a4a5
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 7 deletions.
diff --git a/vulnerability/fetchers/nvd.go b/vulnerability/fetchers/nvd.go
@@ -17,10 +17,12 @@ type NVDFetcher struct {
 	NVDClient NVDClient
 }
 
-func NewNVDFetcher(apiKey string) *NVDFetcher {
-	return &NVDFetcher{
-		NVDKey: apiKey,
-	}
+func NewNVDFetcher() *NVDFetcher {
+	return &NVDFetcher{}
+}
+
+func (fetcher *NVDFetcher) SetNVDKey(key string) {
+	fetcher.NVDKey = key
 }
 
 func (fetcher *NVDFetcher) SetNVDClient(client NVDClient) {

diff --git a/vulnerability/fetchers/nvd_test.go b/vulnerability/fetchers/nvd_test.go
@@ -38,7 +38,8 @@ func (m *MockNVDClientWrapper) ConvertNvdDataToCVE(nvdData fetchers.NVDData) []m
 
 func TestFetchCVEData(t *testing.T) {
 	mockClientWrapper := &MockNVDClientWrapper{}
-	nvdFetcher := fetchers.NewNVDFetcher("mock-api-key")
+	nvdFetcher := fetchers.NewNVDFetcher()
+	nvdFetcher.SetNVDKey("mock-api-key")
 	nvdFetcher.SetNVDClient(mockClientWrapper)
 
 	// Mock the FetchCVEData function in the mock client wrapper.
@@ -69,7 +70,8 @@ func TestFetchCVEData(t *testing.T) {
 
 func TestConvertNvdDataToCVE(t *testing.T) {
 	mockClientWrapper := &MockNVDClientWrapper{}
-	nvdFetcher := fetchers.NewNVDFetcher("mock-api-key")
+	nvdFetcher := fetchers.NewNVDFetcher()
+	nvdFetcher.SetNVDKey("mock-api-key")
 	nvdFetcher.SetNVDClient(mockClientWrapper)
 
 	// Mocking an NVDData for testing purposes.
@@ -144,3 +146,22 @@ func TestParseTime(t *testing.T) {
 		}
 	})
 }
+
+func TestNVDFetcherReal(t *testing.T) {
+	// Create an instance of NVDFetcher for testing.
+	fetcher := fetchers.NewNVDFetcher()
+
+	// Test the FetchCVEData and ConvertNvdDataToCVE function.
+	t.Run("TestNVDFetcherReal", func(t *testing.T) {
+		data, err := fetcher.FetchCVEData()
+		if err != nil {
+			t.Errorf("Error: %v", err)
+		}
+		cves := fetcher.ConvertNvdDataToCVE(data)
+		if len(cves) == 0 {
+			t.Errorf("Error: CVE data not converted")
+		}
+	})
+
+	// You can test other functions here.
+}
diff --git a/vulnerability/service/service.go b/vulnerability/service/service.go
@@ -134,7 +134,8 @@ func (s *VulnService) UpdateCVE(ctx context.Context, req *pb.UpdateCVERequest) (
 }
 
 func (s *VulnService) FetchNVDFeeds(req *pb.FetchNVDFeedsRequest, stream pb.VulnService_FetchNVDFeedsServer) error {
-	fetcher := fetchers.NewNVDFetcher(req.GetApiKey())
+	fetcher := fetchers.NewNVDFetcher()
+	fetcher.SetNVDKey(req.GetApiKey())
 	data, err := fetcher.FetchCVEData()
 	if err != nil {
 		return util.ErrNotPerformedOperation