Skip to content

Commit

Permalink
Fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@Tharsanan1
Tharsanan1 committed Sep 25, 2024
1 parent dfd8b43 commit 634664c
Show file tree
Hide file tree
Showing 4 changed files with 36 additions and 35 deletions.
22 changes: 11 additions & 11 deletions apim-apk-agent/internal/k8sClient/k8s_client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -434,22 +434,22 @@ func DeployAIRateLimitPolicyCR(policy eventhubTypes.SubscriptionPolicy, k8sClien
conf, _ := config.ReadConfigs()
tokenCount := &dpv1alpha3.TokenCount{}
requestCount := &dpv1alpha3.RequestCount{}
if policy.DefaultLimit.AiApiQuota.PromptTokenCount != nil &&
policy.DefaultLimit.AiApiQuota.CompletionTokenCount != nil &&
policy.DefaultLimit.AiApiQuota.TotalTokenCount != nil {
if policy.DefaultLimit.AiAPIQuota.PromptTokenCount != nil &&
policy.DefaultLimit.AiAPIQuota.CompletionTokenCount != nil &&
policy.DefaultLimit.AiAPIQuota.TotalTokenCount != nil {
tokenCount = &dpv1alpha3.TokenCount{
Unit: policy.DefaultLimit.AiApiQuota.TimeUnit,
RequestTokenCount: uint32(*policy.DefaultLimit.AiApiQuota.PromptTokenCount),
ResponseTokenCount: uint32(*policy.DefaultLimit.AiApiQuota.CompletionTokenCount),
TotalTokenCount: uint32(*policy.DefaultLimit.AiApiQuota.TotalTokenCount),
Unit: policy.DefaultLimit.AiAPIQuota.TimeUnit,
RequestTokenCount: uint32(*policy.DefaultLimit.AiAPIQuota.PromptTokenCount),
ResponseTokenCount: uint32(*policy.DefaultLimit.AiAPIQuota.CompletionTokenCount),
TotalTokenCount: uint32(*policy.DefaultLimit.AiAPIQuota.TotalTokenCount),
}
} else {
tokenCount = nil
}
if policy.DefaultLimit.AiApiQuota.RequestCount != nil {
if policy.DefaultLimit.AiAPIQuota.RequestCount != nil {
requestCount = &dpv1alpha3.RequestCount{
RequestsPerUnit: uint32(*policy.DefaultLimit.AiApiQuota.RequestCount),
Unit: policy.DefaultLimit.AiApiQuota.TimeUnit,
RequestsPerUnit: uint32(*policy.DefaultLimit.AiAPIQuota.RequestCount),
Unit: policy.DefaultLimit.AiAPIQuota.TimeUnit,
}
} else {
requestCount = nil
Expand All @@ -462,7 +462,7 @@ func DeployAIRateLimitPolicyCR(policy eventhubTypes.SubscriptionPolicy, k8sClien
Spec: dpv1alpha3.AIRateLimitPolicySpec{
Override: &dpv1alpha3.AIRateLimit{
Organization: policy.TenantDomain,
TokenCount: tokenCount,
TokenCount: tokenCount,
RequestCount: requestCount,
},
TargetRef: gwapiv1b1.PolicyTargetReference{Group: constants.GatewayGroup, Kind: "Subscription", Name: "default"},
Expand Down
28 changes: 14 additions & 14 deletions apim-apk-agent/internal/synchronizer/ratelimit_policy_fetcher.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -238,7 +238,7 @@ func FetchSubscriptionRateLimitPoliciesOnEvent(ratelimitName string, organizatio
logger.LoggerSynchronizer.Infof("Error formatting JSON:", err)
return
}
logger.LoggerSynchronizer.Infof("Json string: %s",string(jsonOutput))
logger.LoggerSynchronizer.Infof("Json string: %s", string(jsonOutput))
logger.LoggerSynchronizer.Debugf("Response String received for Policies: %v", string(responseBytes))

if err != nil {
Expand All @@ -258,27 +258,27 @@ func FetchSubscriptionRateLimitPoliciesOnEvent(ratelimitName string, organizatio
var rateLimitPolicies []eventhubTypes.SubscriptionPolicy = rateLimitPolicyList.List
for _, policy := range rateLimitPolicies {
if policy.QuotaType == "aiApiQuota" {
if policy.DefaultLimit.AiApiQuota != nil {
switch policy.DefaultLimit.AiApiQuota.TimeUnit {
if policy.DefaultLimit.AiAPIQuota != nil {
switch policy.DefaultLimit.AiAPIQuota.TimeUnit {
case "min":
policy.DefaultLimit.AiApiQuota.TimeUnit = "Minute"
policy.DefaultLimit.AiAPIQuota.TimeUnit = "Minute"
case "hours":
policy.DefaultLimit.AiApiQuota.TimeUnit = "Hour"
policy.DefaultLimit.AiAPIQuota.TimeUnit = "Hour"
case "days":
policy.DefaultLimit.AiApiQuota.TimeUnit = "Day"
policy.DefaultLimit.AiAPIQuota.TimeUnit = "Day"
default:
logger.LoggerSynchronizer.Errorf("Unsupported timeunit %s", policy.DefaultLimit.AiApiQuota.TimeUnit)
logger.LoggerSynchronizer.Errorf("Unsupported timeunit %s", policy.DefaultLimit.AiAPIQuota.TimeUnit)
continue
}
if policy.DefaultLimit.AiApiQuota.PromptTokenCount == nil && policy.DefaultLimit.AiApiQuota.TotalTokenCount != nil {
policy.DefaultLimit.AiApiQuota.PromptTokenCount = policy.DefaultLimit.AiApiQuota.TotalTokenCount
if policy.DefaultLimit.AiAPIQuota.PromptTokenCount == nil && policy.DefaultLimit.AiAPIQuota.TotalTokenCount != nil {
policy.DefaultLimit.AiAPIQuota.PromptTokenCount = policy.DefaultLimit.AiAPIQuota.TotalTokenCount
}
if policy.DefaultLimit.AiApiQuota.CompletionTokenCount == nil && policy.DefaultLimit.AiApiQuota.TotalTokenCount != nil {
policy.DefaultLimit.AiApiQuota.CompletionTokenCount = policy.DefaultLimit.AiApiQuota.TotalTokenCount
if policy.DefaultLimit.AiAPIQuota.CompletionTokenCount == nil && policy.DefaultLimit.AiAPIQuota.TotalTokenCount != nil {
policy.DefaultLimit.AiAPIQuota.CompletionTokenCount = policy.DefaultLimit.AiAPIQuota.TotalTokenCount
}
if policy.DefaultLimit.AiApiQuota.TotalTokenCount == nil && policy.DefaultLimit.AiApiQuota.PromptTokenCount != nil && policy.DefaultLimit.AiApiQuota.CompletionTokenCount != nil {
total := *policy.DefaultLimit.AiApiQuota.PromptTokenCount + *policy.DefaultLimit.AiApiQuota.CompletionTokenCount
policy.DefaultLimit.AiApiQuota.TotalTokenCount = &total
if policy.DefaultLimit.AiAPIQuota.TotalTokenCount == nil && policy.DefaultLimit.AiAPIQuota.PromptTokenCount != nil && policy.DefaultLimit.AiAPIQuota.CompletionTokenCount != nil {
total := *policy.DefaultLimit.AiAPIQuota.PromptTokenCount + *policy.DefaultLimit.AiAPIQuota.CompletionTokenCount
policy.DefaultLimit.AiAPIQuota.TotalTokenCount = &total
}
managementserver.AddSubscriptionPolicy(policy)
k8sclient.DeployAIRateLimitPolicyCR(policy, c)
Expand Down
15 changes: 8 additions & 7 deletions apim-apk-agent/pkg/eventhub/types/types.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -196,8 +196,8 @@ type ConditionGroup struct {

// DefaultLimit represents the default limit within the response.
type DefaultLimit struct {
AiApiQuota *AiApiQuota `json:"aiApiQuota"`
QuotaType string `json:"quotaType"`
AiAPIQuota *AiAPIQuota `json:"aiApiQuota"`
QuotaType string `json:"quotaType"`
RequestCount struct {
TimeUnit string `json:"timeUnit"`
UnitTime int `json:"unitTime"`
Expand All @@ -207,12 +207,13 @@ type DefaultLimit struct {
EventCount interface{} `json:"eventCount"`
}

type AiApiQuota struct {
CompletionTokenCount *int `json:"completionTokenCount"`
PromptTokenCount *int `json:"promptTokenCount"`
RequestCount *int `json:"requestCount"`
// AiAPIQuota contains the AI ratelimit configurations
type AiAPIQuota struct {
CompletionTokenCount *int `json:"completionTokenCount"`
PromptTokenCount *int `json:"promptTokenCount"`
RequestCount *int `json:"requestCount"`
TimeUnit string `json:"timeUnit"`
TotalTokenCount *int `json:"totalTokenCount"`
TotalTokenCount *int `json:"totalTokenCount"`
UnitTime int `json:"unitTime"`
}

Expand Down
6 changes: 3 additions & 3 deletions apim-apk-agent/pkg/transformer/transformer_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import (
"testing"

"github.com/stretchr/testify/assert"
dpv1alpha2 "github.com/wso2/apk/common-go-libs/apis/dp/v1alpha2"
dpv1alpha3 "github.com/wso2/apk/common-go-libs/apis/dp/v1alpha3"
gwapiv1 "sigs.k8s.io/gateway-api/apis/v1"
)

Expand Down Expand Up @@ -307,7 +307,7 @@ func TestReplaceVhost(t *testing.T) {
assert.Equal(t, gwapiv1.Hostname(env.Vhost), k8sArtifact.HTTPRoutes[routes].Spec.Hostnames[0])
}
}
assert.IsType(t, []dpv1alpha2.EnvConfig{}, k8sArtifact.API.Spec.Production)
assert.IsType(t, []dpv1alpha3.EnvConfig{}, k8sArtifact.API.Spec.Production)
assert.Empty(t, k8sArtifact.API.Spec.Production, "Production should be empty")

} else {
Expand All @@ -316,7 +316,7 @@ func TestReplaceVhost(t *testing.T) {
assert.Equal(t, gwapiv1.Hostname(env.Vhost), k8sArtifact.HTTPRoutes[routes].Spec.Hostnames[0])
}
}
assert.IsType(t, []dpv1alpha2.EnvConfig{}, k8sArtifact.API.Spec.Sandbox)
assert.IsType(t, []dpv1alpha3.EnvConfig{}, k8sArtifact.API.Spec.Sandbox)
assert.Empty(t, k8sArtifact.API.Spec.Sandbox, "Sandbox should be empty")
}
}
Expand Down

0 comments on commit 634664c

Please sign in to comment.