Add integration tests for AI API related import export flows

import-export-cli/integration/api_test.go

@@ -1448,6 +1448,29 @@ func TestExportImportWebSocketApiFromAsyncApiDef(t *testing.T) {
 	}
 }
 
+// Export an AI API from one environment and import to another environment by specifying the provider name
+func TestExportImportAiApi(t *testing.T) {
+	for _, user := range testCaseUsers {
+		t.Run(user.Description, func(t *testing.T) {
+			dev := GetDevClient()
+			prod := GetProdClient()
+
+			api := testutils.AddAIAPIFromOpenAPIDefinition(t, dev, user.ApiCreator.Username, user.ApiCreator.Password)
+
+			args := &testutils.ApiImportExportTestArgs{
+				ApiProvider: testutils.Credentials{Username: user.ApiCreator.Username, Password: user.ApiCreator.Password},
+				CtlUser:     testutils.Credentials{Username: user.CtlUser.Username, Password: user.CtlUser.Password},
+				Api:         api,
+				SrcAPIM:     dev,
+				DestAPIM:    prod,
+			}
+
+			testutils.ValidateAIAPIExportImport(t, args, testutils.APITypeAI)
+
+		})
+	}
+}
+
 // Import an API and then create a new version of that API by updating the context and version only and import again
 func TestCreateNewVersionOfApiByUpdatingVersion(t *testing.T) {
 	for _, user := range testCaseUsers {

import-export-cli/integration/apim/apim.go

@@ -44,10 +44,11 @@ import (
 const (
 	restClientPostfix = "-integ-rest-client"
 
-	ApplicationThrottlePolicyType  = "application"
-	CustomThrottlePolicyType       = "custom"
-	AdvancedThrottlePolicyType     = "advanced"
-	SubscriptionThrottlePolicyType = "subscription"
+	ApplicationThrottlePolicyType       = "application"
+	CustomThrottlePolicyType            = "custom"
+	AdvancedThrottlePolicyType          = "advanced"
+	SubscriptionThrottlePolicyType      = "subscription"
+	AiApiSubscriptionThrottlePolicyType = "aiApiSubscription"
 
 	applicationThrottlePolicyQuery  = "app"
 	customThrottlePolicyQuery       = "global"
@@ -285,6 +286,27 @@ func (instance *Client) GenerateAdditionalProperties(provider, endpointUrl, apiT
 		advertiseInfo, _ := json.Marshal(api.AdvertiseInformation)
 		additionalProperties = additionalProperties + `"type":"` + apiType + `",
 		"advertiseInfo": ` + string(advertiseInfo) + `}`
+	} else if strings.EqualFold(apiType, "AIAPI") {
+		additionalProperties = additionalProperties +
+			`"subtypeConfiguration": {
+				"subtype":"AIAPI",
+				"configuration":{
+					"llmProviderName": "MistralAI",
+					"llmProviderApiVersion": "1.0.0"
+				}
+			},
+			"securityScheme": ["api_key"],
+			"egress": true,
+			"endpointConfig": {   
+				"endpoint_type":"http",
+				"sandbox_endpoints":{
+					"url":"` + endpointUrl + `"
+				},
+				"production_endpoints":{
+					"url":"` + endpointUrl + `"
+				}
+			}
+		}`
 	} else {
 		additionalProperties = additionalProperties +
 			`"endpointConfig": {   
@@ -682,6 +704,56 @@ func (instance *Client) AddAPIFromOpenAPIDefinition(t *testing.T, path string, a
 	return apiResponse.ID
 }
 
+// AddAIAPIFromOpenAPIDefinition : Add Mistral AI API using an OpenAPI Definition to APIM
+func (instance *Client) AddAIAPIFromOpenAPIDefinition(t *testing.T, path string, additionalProperties string, username string, password string) string {
+	apisURL := instance.publisherRestURL + "/apis/import-openapi"
+
+	inlineAPIDefinition, err := ioutil.ReadFile(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	body := &bytes.Buffer{}
+	writer := multipart.NewWriter(body)
+
+	part, err := writer.CreateFormField("inlineAPIDefinition")
+	if err != nil {
+		t.Fatal(err)
+	}
+	part.Write([]byte(inlineAPIDefinition))
+
+	part, err = writer.CreateFormField("additionalProperties")
+	if err != nil {
+		t.Fatal(err)
+	}
+	part.Write([]byte(additionalProperties))
+
+	err = writer.Close()
+
+	request := base.CreatePost(apisURL, body)
+
+	base.SetDefaultRestAPIHeadersToConsumeFormData(instance.accessToken, request)
+
+	base.LogRequest("apim.AddAIAPIFromOpenAPIDefinition()", request)
+
+	response := base.SendHTTPRequest(request)
+
+	defer response.Body.Close()
+
+	base.ValidateAndLogResponse("apim.AddAIAPIFromOpenAPIDefinition()", response, 201)
+
+	var apiResponse API
+	json.NewDecoder(response.Body).Decode(&apiResponse)
+
+	t.Cleanup(func() {
+		username, password := RetrieveAdminCredentialsInsteadCreator(username, password)
+		instance.Login(username, password)
+		instance.DeleteAPI(apiResponse.ID)
+	})
+
+	return apiResponse.ID
+}
+
 // AddGraphQLAPI : Add new GraphQL API to APIM
 func (instance *Client) AddGraphQLAPI(t *testing.T, path, additionalProperties, username, password string) string {
 	apisURL := instance.publisherRestURL + "/apis/import-graphql-schema"
@@ -2111,6 +2183,11 @@ func (instance *Client) registerClient(username string, password string) dcrResp
 func (instance *Client) AddThrottlePolicy(t *testing.T, policy interface{}, username, password, policyType string, doClean bool) map[string]interface{} {
 	var throttlePolicyResponse map[string]interface{}
 
+	// AiApiSubscriptionThrottlePolicyType should have the type as SubscriptionThrottlePolicyType as it falls under the type "sub"
+	if policyType == AiApiSubscriptionThrottlePolicyType {
+		policyType = SubscriptionThrottlePolicyType
+	}
+
 	throttlePolicyURL := instance.adminRestURL + "/throttling/policies/" + policyType
 
 	data, err := json.Marshal(policy)
@@ -2249,11 +2326,15 @@ func (instance *Client) GetThrottlePolicies(t *testing.T) *utils.ThrottlingPolic
 // GenerateSampleThrottlePolicyData : Generate sample ThrottlePolicy of a specific throttling policy type
 func (instance *Client) GenerateSampleThrottlePolicyData(policyType string) interface{} {
 	const (
-		policyString      = "Policy"
-		requestCountLimit = "REQUESTCOUNTLIMIT"
-		timeUnit          = "min"
-		unitTime          = 10
-		requestCount      = 5
+		policyString         = "Policy"
+		requestCountLimit    = "REQUESTCOUNTLIMIT"
+		aiApiQuotaLimit      = "AIAPIQUOTALIMIT"
+		timeUnit             = "min"
+		unitTime             = 10
+		requestCount         = 5
+		totalTokenCount      = 50
+		promptTokenCount     = 20
+		completionTokenCount = 30
 	)
 
 	switch policyType {
@@ -2317,6 +2398,22 @@ func (instance *Client) GenerateSampleThrottlePolicyData(policyType string) inte
 		policy.Permissions.PermissionType = "ALLOW"
 		policy.Permissions.Roles = []string{"admin"}
 		return &policy
+	case AiApiSubscriptionThrottlePolicyType:
+		policy := SubscriptionThrottlePolicy{}
+		policy.PolicyName = base.GenerateRandomString() + policyString
+		policy.Description = "This is a Test AI API Subscription Policy"
+		policy.IsDeployed = false
+		policy.Type = "SubscriptionThrottlePolicy"
+		policy.DefaultLimit.Type = aiApiQuotaLimit
+		policy.DefaultLimit.AiApiQuota.TimeUnit = timeUnit
+		policy.DefaultLimit.AiApiQuota.UnitTime = unitTime
+		policy.DefaultLimit.AiApiQuota.RequestCount = requestCount
+		policy.DefaultLimit.AiApiQuota.TotalTokenCount = totalTokenCount
+		policy.DefaultLimit.AiApiQuota.PromptTokenCount = promptTokenCount
+		policy.DefaultLimit.AiApiQuota.CompletionTokenCount = completionTokenCount
+		policy.Permissions.PermissionType = "ALLOW"
+		policy.Permissions.Roles = []string{"admin"}
+		return &policy
 	}
 	return nil
 }

import-export-cli/integration/apim/throttlePolicyDTO.go

@@ -112,6 +112,7 @@ type DefaultLimit struct {
 	RequestCount RequestCount `json:"requestCount"`
 	Bandwidth    Bandwidth    `json:"bandwidth"`
 	EventCount   EventCount   `json:"eventCount"`
+	AiApiQuota   AiApiQuota   `json:"aiApiQuota"`
 }
 
 type RequestCount struct {
@@ -132,3 +133,12 @@ type EventCount struct {
 	UnitTime   int    `json:"unitTime"`
 	EventCount int    `json:"eventCount"`
 }
+
+type AiApiQuota struct {
+	TimeUnit             string `json:"timeUnit"`
+	UnitTime             int    `json:"unitTime"`
+	RequestCount         int    `json:"requestCount"`
+	TotalTokenCount      int    `json:"totalTokenCount"`
+	PromptTokenCount     int    `json:"promptTokenCount"`
+	CompletionTokenCount int    `json:"completionTokenCount"`
+}