From bb7a4a55288f23e4ae4ee87f68ccfafc51429dbe Mon Sep 17 00:00:00 2001
From: mtnmunuklu <mtnmunuklu@gmail.com>
Date: Tue, 13 Feb 2024 22:16:35 +0300
Subject: [PATCH] Modify NewNVDFetcher to allow setting NVDKey separately

---
 vulnerability/fetchers/nvd.go      | 10 ++++++----
 vulnerability/fetchers/nvd_test.go | 25 +++++++++++++++++++++++--
 vulnerability/service/service.go   |  3 ++-
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/vulnerability/fetchers/nvd.go b/vulnerability/fetchers/nvd.go
index b7ae165..c4ff3f4 100644
--- 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
index 8890d71..494bec1 100644
--- 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
index 11a1d4e..c4fbb6f 100644
--- 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