Merge branch 'main' into sofa-url-column

macadmins · May 30, 2024 · 5861548 · 5861548
2 parents 7e69d3d + 1c2c38c
commit 5861548
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 7 deletions.
diff --git a/tables/sofa/client_test.go b/tables/sofa/client_test.go
@@ -118,11 +118,16 @@ func TestUnmarshalJSON(t *testing.T) {
 }
 
 func TestWithUserAgent(t *testing.T) {
-	cwd, err := os.Getwd()
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "Foo/2.0", r.Header.Get("User-Agent"))
+	}))
+	path := t.TempDir()
+	client, err := NewSofaClient(WithUserAgent("Foo/2.0"), WithCacheDir(path))
 	assert.NoError(t, err)
-	client, err := NewSofaClient(WithUserAgent("test"), WithCacheDir(cwd))
+	assert.Equal(t, "Foo/2.0", client.userAgent)
+
+	err = client.downloadFile(server.URL, path+"/test.txt")
 	assert.NoError(t, err)
-	assert.Equal(t, "test", client.userAgent)
 }
 
 func TestLoacCachedEtag(t *testing.T) {

diff --git a/tables/sofa/sofa_cves.go b/tables/sofa/sofa_cves.go
@@ -26,7 +26,7 @@ func SofaUnpatchedCVEsColumns() []table.ColumnDefinition {
 	}
 }
 
-func SofaUnpatchedCVEsGenerate(ctx context.Context, queryContext table.QueryContext, socketPath string) ([]map[string]string, error) {
+func SofaUnpatchedCVEsGenerate(ctx context.Context, queryContext table.QueryContext, socketPath string, opts ...Option) ([]map[string]string, error) {
 	url := SofaV1URL
 	if constraintList, present := queryContext.Constraints["url"]; present {
 		// 'url' is in the where clause
@@ -62,7 +62,12 @@ func SofaUnpatchedCVEsGenerate(ctx context.Context, queryContext table.QueryCont
 		}
 	}
 
-	client, err := NewSofaClient(WithURL(url))
+	defaultOpts := []Option{
+		WithURL(url),
+	}
+	opts = append(opts, defaultOpts...)
+
+	client, err := NewSofaClient(opts...)
 	if err != nil {
 		return nil, err
 	}

diff --git a/tables/sofa/sofa_info.go b/tables/sofa/sofa_info.go
@@ -131,7 +131,7 @@ func SofaSecurityReleaseInfoColumns() []table.ColumnDefinition {
 	}
 }
 
-func SofaSecurityReleaseInfoGenerate(ctx context.Context, queryContext table.QueryContext, socketPath string) ([]map[string]string, error) {
+func SofaSecurityReleaseInfoGenerate(ctx context.Context, queryContext table.QueryContext, socketPath string, clientOpts ...Option) ([]map[string]string, error) {
 	url, osVersion := processContextConstraints(queryContext)
 
 	if osVersion == "" {
@@ -148,7 +148,12 @@ func SofaSecurityReleaseInfoGenerate(ctx context.Context, queryContext table.Que
 		}
 	}
 
-	client, err := NewSofaClient(WithURL(url))
+	defaultOpts := []Option{
+		WithURL(url),
+	}
+	clientOpts = append(clientOpts, defaultOpts...)
+
+	client, err := NewSofaClient(clientOpts...)
 	if err != nil {
 		return nil, err
 	}