feat: associate team

.github/workflows/test-upload-file.yml

@@ -92,3 +92,11 @@ jobs:
           security_agent_api_key: ${{ secrets.SECURITY_AGENT_API_KEY }}
           tags: |
             product_release=v1.0.0
+
+      - name: Test team association
+        uses: ./upload-file
+        with:
+          file_path: test-upload.json
+          security_agent_api_key: ${{ secrets.SECURITY_AGENT_API_KEY }}
+          tags: |
+            team=security

upload-file/catalog/teams.go

@@ -0,0 +1,34 @@
+package catalog
+
+import (
+	"context"
+
+	"github.com/hasura/security-agent-tools/upload-file/saclient"
+	"github.com/machinebox/graphql"
+)
+
+func Teams(ctx context.Context, c *saclient.Client) ([]string, error) {
+	req := graphql.NewRequest(`query GetTeams {
+  team_catalog_teams {
+    name
+  }
+}`)
+
+	var response struct {
+		Teams []struct {
+			Name string `json:"name"`
+		} `json:"team_catalog_teams"`
+	}
+
+	err := c.ExecuteGQL(ctx, req, &response)
+	if err != nil {
+		return nil, err
+	}
+
+	var teams []string
+	for _, team := range response.Teams {
+		teams = append(teams, team.Name)
+	}
+
+	return teams, nil
+}

upload-file/catalog/teams_test.go

@@ -0,0 +1,155 @@
+package catalog
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/hasura/security-agent-tools/upload-file/saclient"
+)
+
+func TestTeams_Success(t *testing.T) {
+	expectedTeams := []string{"backend-team", "frontend-team", "devops-team"}
+
+	// Mock GraphQL server
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Verify the request contains the expected query
+		body := make([]byte, r.ContentLength)
+		r.Body.Read(body)
+		if !strings.Contains(string(body), "team_catalog_teams") {
+			t.Error("Expected query to contain team_catalog_teams")
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{
+			"data": {
+				"team_catalog_teams": [
+					{"name": "backend-team"},
+					{"name": "frontend-team"},
+					{"name": "devops-team"}
+				]
+			}
+		}`))
+	}))
+	defer server.Close()
+
+	client := saclient.NewClient(server.URL, "test-api-key")
+	teams, err := Teams(context.Background(), client)
+
+	if err != nil {
+		t.Errorf("Teams failed: %v", err)
+	}
+
+	if len(teams) != len(expectedTeams) {
+		t.Errorf("Expected %d teams, got %d", len(expectedTeams), len(teams))
+	}
+
+	for i, expectedTeam := range expectedTeams {
+		if i >= len(teams) || teams[i] != expectedTeam {
+			t.Errorf("Expected team %s at index %d, got %s", expectedTeam, i, teams[i])
+		}
+	}
+}
+
+func TestTeams_EmptyResponse(t *testing.T) {
+	// Mock GraphQL server that returns empty array
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{
+			"data": {
+				"team_catalog_teams": []
+			}
+		}`))
+	}))
+	defer server.Close()
+
+	client := saclient.NewClient(server.URL, "test-api-key")
+	teams, err := Teams(context.Background(), client)
+
+	if err != nil {
+		t.Errorf("Teams failed: %v", err)
+	}
+
+	if len(teams) != 0 {
+		t.Errorf("Expected empty teams array, got %d teams", len(teams))
+	}
+}
+
+func TestTeams_SingleTeam(t *testing.T) {
+	expectedTeam := "security-team"
+
+	// Mock GraphQL server
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{
+			"data": {
+				"team_catalog_teams": [
+					{"name": "security-team"}
+				]
+			}
+		}`))
+	}))
+	defer server.Close()
+
+	client := saclient.NewClient(server.URL, "test-api-key")
+	teams, err := Teams(context.Background(), client)
+
+	if err != nil {
+		t.Errorf("Teams failed: %v", err)
+	}
+
+	if len(teams) != 1 {
+		t.Errorf("Expected 1 team, got %d", len(teams))
+	}
+
+	if teams[0] != expectedTeam {
+		t.Errorf("Expected team %s, got %s", expectedTeam, teams[0])
+	}
+}
+
+func TestTeams_GraphQLError(t *testing.T) {
+	// Mock GraphQL server that returns an error
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{"errors": [{"message": "Database error"}]}`))
+	}))
+	defer server.Close()
+
+	client := saclient.NewClient(server.URL, "test-api-key")
+	teams, err := Teams(context.Background(), client)
+
+	if err == nil {
+		t.Error("Expected error for GraphQL error response")
+	}
+
+	if teams != nil {
+		t.Error("Expected teams to be nil when error occurs")
+	}
+}
+
+func TestTeams_MalformedJSON(t *testing.T) {
+	// Mock GraphQL server that returns malformed JSON
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{"data": {"team_catalog_teams": [{"name": "incomplete"`))
+	}))
+	defer server.Close()
+
+	client := saclient.NewClient(server.URL, "test-api-key")
+	teams, err := Teams(context.Background(), client)
+
+	if err == nil {
+		t.Error("Expected error for malformed JSON response")
+	}
+
+	if teams != nil {
+		t.Error("Expected teams to be nil when JSON parsing fails")
+	}
+}

upload-file/main.go

@@ -89,4 +89,17 @@ func main() {
 	case productRelease != "":
 		log.Printf("Associated product release: %s\n", productRelease)
 	}
+
+	team, err := sc.AssociateTeam(context.Background())
+	switch {
+	case err != nil:
+		teams, _ := catalog.Teams(context.Background(), secAgentClient)
+		var ts strings.Builder
+		for _, t := range teams {
+			ts.WriteString("  - " + t + "\n")
+		}
+		log.Fatalf("Failed to associate team with scan: %v. Please check `team` value is one of the following:\n%s", err, ts.String())
+	case team != "":
+		log.Printf("Associated team: %s\n", team)
+	}
 }

upload-file/scan/associate_team.go

@@ -0,0 +1,36 @@
+package scan
+
+import (
+	"context"
+
+	"github.com/machinebox/graphql"
+)
+
+func (s *Scan) AssociateTeam(ctx context.Context) (string, error) {
+	team := s.Tags["team"]
+	if team == "" {
+		return "", nil
+	}
+
+	req := graphql.NewRequest(`mutation AssociateTeam($scan_id: uuid!, $team_name: string!) {
+  insert_vulnerability_reports_by_team(
+    objects: {scan_id: $scan_id, team_name: $team_name}
+  ) {
+    returning {
+      id
+    }
+  }
+}`)
+	req.Var("scan_id", s.ID)
+	req.Var("team_name", team)
+
+	var response struct {
+		InsertVulnerabilityReportsByTeam struct {
+			Returning []struct {
+				ID string `json:"id"`
+			} `json:"returning"`
+		} `json:"insert_vulnerability_reports_by_team"`
+	}
+
+	return team, s.client.ExecuteGQL(ctx, req, &response)
+}

upload-file/scan/associate_team_test.go

@@ -0,0 +1,97 @@
+package scan
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/hasura/security-agent-tools/upload-file/saclient"
+)
+
+func TestScan_AssociateTeam_Success(t *testing.T) {
+	expectedTeam := "backend-team"
+
+	// Mock GraphQL server
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{
+			"data": {
+				"insert_vulnerability_reports_by_team": {
+					"returning": [{
+						"id": "team-assoc-id-456"
+					}]
+				}
+			}
+		}`))
+	}))
+	defer server.Close()
+
+	client := saclient.NewClient(server.URL, "test-api-key")
+	scan := &Scan{
+		ID: "scan-id-123",
+		Tags: map[string]string{
+			"team": expectedTeam,
+		},
+		client: client,
+	}
+
+	team, err := scan.AssociateTeam(context.Background())
+
+	if err != nil {
+		t.Errorf("AssociateTeam failed: %v", err)
+	}
+
+	if team != expectedTeam {
+		t.Errorf("Expected team %s, got %s", expectedTeam, team)
+	}
+}
+
+func TestScan_AssociateTeam_EmptyTeam(t *testing.T) {
+	client := saclient.NewClient("https://example.com/graphql", "test-api-key")
+	scan := &Scan{
+		ID:     "scan-id-123",
+		Tags:   map[string]string{}, // No team tag
+		client: client,
+	}
+
+	team, err := scan.AssociateTeam(context.Background())
+
+	if err != nil {
+		t.Errorf("AssociateTeam failed: %v", err)
+	}
+
+	if team != "" {
+		t.Errorf("Expected empty team, got %s", team)
+	}
+}
+
+func TestScan_AssociateTeam_GraphQLError(t *testing.T) {
+	// Mock GraphQL server that returns an error
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{"errors": [{"message": "Database error"}]}`))
+	}))
+	defer server.Close()
+
+	client := saclient.NewClient(server.URL, "test-api-key")
+	scan := &Scan{
+		ID: "scan-id-123",
+		Tags: map[string]string{
+			"team": "backend-team",
+		},
+		client: client,
+	}
+
+	team, err := scan.AssociateTeam(context.Background())
+
+	if err == nil {
+		t.Error("Expected error for GraphQL error response")
+	}
+
+	if team != "backend-team" {
+		t.Errorf("Expected team to be returned even on error, got %s", team)
+	}
+}