Add code_search and dependency_snapshots for RateLimits (#3019)

Fixes: #3018.

github/github.go

@@ -1303,6 +1303,8 @@ const (
 	codeScanningUploadCategory
 	actionsRunnerRegistrationCategory
 	scimCategory
+	dependencySnapshotsCategory
+	codeSearchCategory
 
 	categories // An array of this length will be able to contain all rate limit categories.
 )
@@ -1315,6 +1317,12 @@ func category(method, path string) rateLimitCategory {
 		// NOTE: coreCategory is returned for actionsRunnerRegistrationCategory too,
 		// because no API found for this category.
 		return coreCategory
+
+	// https://docs.github.com/en/rest/search/search#search-code
+	case strings.HasPrefix(path, "/search/code") &&
+		method == http.MethodGet:
+		return codeSearchCategory
+
 	case strings.HasPrefix(path, "/search/"):
 		return searchCategory
 	case path == "/graphql":
@@ -1337,6 +1345,12 @@ func category(method, path string) rateLimitCategory {
 	// https://docs.github.com/enterprise-cloud@latest/rest/scim
 	case strings.HasPrefix(path, "/scim/"):
 		return scimCategory
+
+	// https://docs.github.com/en/rest/dependency-graph/dependency-submission#create-a-snapshot-of-dependencies-for-a-repository
+	case strings.HasPrefix(path, "/repos/") &&
+		strings.HasSuffix(path, "/dependency-graph/snapshots") &&
+		method == http.MethodPost:
+		return dependencySnapshotsCategory
 	}
 }
 

github/github_test.go

@@ -1153,6 +1153,16 @@ func TestDo_rateLimitCategory(t *testing.T) {
 			url:      "/scim/v2/organizations/ORG/Users",
 			category: scimCategory,
 		},
+		{
+			method:   http.MethodPost,
+			url:      "/repos/google/go-github/dependency-graph/snapshots",
+			category: dependencySnapshotsCategory,
+		},
+		{
+			method:   http.MethodGet,
+			url:      "/search/code?q=rate",
+			category: codeSearchCategory,
+		},
 		// missing a check for actionsRunnerRegistrationCategory: API not found
 	}
 

github/rate_limit.go

@@ -52,6 +52,8 @@ type RateLimits struct {
 	CodeScanningUpload        *Rate `json:"code_scanning_upload"`
 	ActionsRunnerRegistration *Rate `json:"actions_runner_registration"`
 	SCIM                      *Rate `json:"scim"`
+	DependencySnapshots       *Rate `json:"dependency_snapshots"`
+	CodeSearch                *Rate `json:"code_search"`
 }
 
 func (r RateLimits) String() string {
@@ -106,6 +108,12 @@ func (s *RateLimitService) Get(ctx context.Context) (*RateLimits, *Response, err
 		if response.Resources.SCIM != nil {
 			s.client.rateLimits[scimCategory] = *response.Resources.SCIM
 		}
+		if response.Resources.DependencySnapshots != nil {
+			s.client.rateLimits[dependencySnapshotsCategory] = *response.Resources.DependencySnapshots
+		}
+		if response.Resources.CodeSearch != nil {
+			s.client.rateLimits[codeSearchCategory] = *response.Resources.CodeSearch
+		}
 		s.client.rateMu.Unlock()
 	}
 

github/rate_limit_test.go

@@ -25,8 +25,10 @@ func TestRateLimits_String(t *testing.T) {
 		CodeScanningUpload:        &Rate{},
 		ActionsRunnerRegistration: &Rate{},
 		SCIM:                      &Rate{},
+		DependencySnapshots:       &Rate{},
+		CodeSearch:                &Rate{},
 	}
-	want := `github.RateLimits{Core:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, Search:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, GraphQL:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, IntegrationManifest:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, SourceImport:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, CodeScanningUpload:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, ActionsRunnerRegistration:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, SCIM:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}}`
+	want := `github.RateLimits{Core:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, Search:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, GraphQL:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, IntegrationManifest:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, SourceImport:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, CodeScanningUpload:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, ActionsRunnerRegistration:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, SCIM:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, DependencySnapshots:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}, CodeSearch:github.Rate{Limit:0, Remaining:0, Reset:github.Timestamp{0001-01-01 00:00:00 +0000 UTC}}}`
 	if got := v.String(); got != want {
 		t.Errorf("RateLimits.String = %v, want %v", got, want)
 	}
@@ -46,7 +48,9 @@ func TestRateLimits(t *testing.T) {
 			"source_import": {"limit":6,"remaining":5,"reset":1372700877},
 			"code_scanning_upload": {"limit":7,"remaining":6,"reset":1372700878},
 			"actions_runner_registration": {"limit":8,"remaining":7,"reset":1372700879},
-			"scim": {"limit":9,"remaining":8,"reset":1372700880}
+			"scim": {"limit":9,"remaining":8,"reset":1372700880},
+			"dependency_snapshots": {"limit":10,"remaining":9,"reset":1372700881},
+			"code_search": {"limit":11,"remaining":10,"reset":1372700882}
 		}}`)
 	})
 
@@ -97,6 +101,16 @@ func TestRateLimits(t *testing.T) {
 			Remaining: 8,
 			Reset:     Timestamp{time.Date(2013, time.July, 1, 17, 48, 00, 0, time.UTC).Local()},
 		},
+		DependencySnapshots: &Rate{
+			Limit:     10,
+			Remaining: 9,
+			Reset:     Timestamp{time.Date(2013, time.July, 1, 17, 48, 1, 0, time.UTC).Local()},
+		},
+		CodeSearch: &Rate{
+			Limit:     11,
+			Remaining: 10,
+			Reset:     Timestamp{time.Date(2013, time.July, 1, 17, 48, 2, 0, time.UTC).Local()},
+		},
 	}
 	if !cmp.Equal(rate, want) {
 		t.Errorf("RateLimits returned %+v, want %+v", rate, want)
@@ -137,6 +151,14 @@ func TestRateLimits(t *testing.T) {
 			category: scimCategory,
 			rate:     want.SCIM,
 		},
+		{
+			category: dependencySnapshotsCategory,
+			rate:     want.DependencySnapshots,
+		},
+		{
+			category: codeSearchCategory,
+			rate:     want.CodeSearch,
+		},
 	}
 
 	for _, tt := range tests {
@@ -177,7 +199,9 @@ func TestRateLimits_overQuota(t *testing.T) {
 			"source_import": {"limit":6,"remaining":5,"reset":1372700877},
 			"code_scanning_upload": {"limit":7,"remaining":6,"reset":1372700878},
 			"actions_runner_registration": {"limit":8,"remaining":7,"reset":1372700879},
-			"scim": {"limit":9,"remaining":8,"reset":1372700880}
+			"scim": {"limit":9,"remaining":8,"reset":1372700880},
+			"dependency_snapshots": {"limit":10,"remaining":9,"reset":1372700881},
+			"code_search": {"limit":11,"remaining":10,"reset":1372700882}
 		}}`)
 	})
 
@@ -228,6 +252,16 @@ func TestRateLimits_overQuota(t *testing.T) {
 			Remaining: 8,
 			Reset:     Timestamp{time.Date(2013, time.July, 1, 17, 48, 00, 0, time.UTC).Local()},
 		},
+		DependencySnapshots: &Rate{
+			Limit:     10,
+			Remaining: 9,
+			Reset:     Timestamp{time.Date(2013, time.July, 1, 17, 48, 1, 0, time.UTC).Local()},
+		},
+		CodeSearch: &Rate{
+			Limit:     11,
+			Remaining: 10,
+			Reset:     Timestamp{time.Date(2013, time.July, 1, 17, 48, 2, 0, time.UTC).Local()},
+		},
 	}
 	if !cmp.Equal(rate, want) {
 		t.Errorf("RateLimits returned %+v, want %+v", rate, want)
@@ -269,6 +303,14 @@ func TestRateLimits_overQuota(t *testing.T) {
 			category: scimCategory,
 			rate:     want.SCIM,
 		},
+		{
+			category: dependencySnapshotsCategory,
+			rate:     want.DependencySnapshots,
+		},
+		{
+			category: codeSearchCategory,
+			rate:     want.CodeSearch,
+		},
 	}
 	for _, tt := range tests {
 		if got, want := client.rateLimits[tt.category], *tt.rate; got != want {
@@ -321,6 +363,16 @@ func TestRateLimits_Marshal(t *testing.T) {
 			Remaining: 1,
 			Reset:     Timestamp{referenceTime},
 		},
+		DependencySnapshots: &Rate{
+			Limit:     1,
+			Remaining: 1,
+			Reset:     Timestamp{referenceTime},
+		},
+		CodeSearch: &Rate{
+			Limit:     1,
+			Remaining: 1,
+			Reset:     Timestamp{referenceTime},
+		},
 	}
 
 	want := `{
@@ -363,6 +415,16 @@ func TestRateLimits_Marshal(t *testing.T) {
 			"limit": 1,
 			"remaining": 1,
 			"reset": ` + referenceTimeStr + `
+		},
+		"dependency_snapshots": {
+			"limit": 1,
+			"remaining": 1,
+			"reset": ` + referenceTimeStr + `
+		},
+		"code_search": {
+			"limit": 1,
+			"remaining": 1,
+			"reset": ` + referenceTimeStr + `
 		}
 	}`