Skip to content

Commit dc3d8bd

Browse files
ashera96ashera96
committed
Fix AI API testcase
1 parent 238d2e9 commit dc3d8bd

File tree

1 file changed

+63
-50
lines changed
  • all-in-one-apim/modules/integration/tests-integration/tests-backend/src/test/java/org/wso2/am/integration/tests/aiapi

1 file changed

+63
-50
lines changed

‎all-in-one-apim/modules/integration/tests-integration/tests-backend/src/test/java/org/wso2/am/integration/tests/aiapi/AIAPITestCase.java‎

Lines changed: 63 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,6 @@ public class AIAPITestCase extends APIMIntegrationBaseTest {
131131
private final String defaultProductionEndpointId = "default_production_endpoint";
132132
private final String defaultProductionEndpointName = "Default Production Endpoint";
133133
private final String newVersion = "2.0.0";
134-
private final String newVersionContext = mistralAPIContext + "-v2";
135134
private final String failoverEndpointUrl = "/failover-endpoint";
136135
private final String failoverEndpointName = "Failover Endpoint";
137136
private String productionEndpointId;
@@ -816,7 +815,7 @@ public void testAIApiInvocationAfterAddingApiLevelModelRoundRobinPolicy() throws
816815
public void testCreateNewVersionWithFailoverPolicy() throws Exception {
817816
// Create new version of the API
818817
HttpResponse copyApiResponse = restAPIPublisher.copyAPI(newVersion, mistralAPIId, false);
819-
Assert.assertEquals(copyApiResponse.getResponseCode(), HttpStatus.SC_CREATED,
818+
Assert.assertEquals(copyApiResponse.getResponseCode(), HttpStatus.SC_OK,
820819
"Failed to create new version of API");
821820

822821
newVersionApiId = copyApiResponse.getData();
@@ -829,38 +828,64 @@ public void testCreateNewVersionWithFailoverPolicy() throws Exception {
829828

830829
APIDTO newVersionApiDto = new Gson().fromJson(getNewVersionApiResponse.getData(), APIDTO.class);
831830
Assert.assertEquals(newVersionApiDto.getVersion(), newVersion, "API version should match");
831+
832+
// Add a new endpoint with failover URL for the new version
833+
String failoverEndpointConfig = readFile(resourcePath + "prod-endpoint-add-endpoint-config.json");
834+
JSONObject failoverEndpointConfigObj = new JSONObject(failoverEndpointConfig);
835+
JSONObject productionEndpoints = new JSONObject();
836+
productionEndpoints.put("url", endpointHost + ":" + endpointPort + failoverEndpointUrl);
837+
failoverEndpointConfigObj.put("production_endpoints", productionEndpoints);
838+
839+
Map<String, Object> failoverEndpointConfigMap = new Gson().fromJson(failoverEndpointConfigObj.toString(), Map.class);
840+
HttpResponse addFailoverEndpointResponse = restAPIPublisher.addApiEndpoint(newVersionApiId, failoverEndpointName,
841+
productionDeploymentStage, failoverEndpointConfigMap);
842+
Assert.assertEquals(addFailoverEndpointResponse.getResponseCode(), HttpStatus.SC_CREATED,
843+
"Failed to add failover endpoint to new version API");
832844

845+
// Extract failover endpoint ID from response
846+
JSONObject failoverEndpointResponse = new JSONObject(addFailoverEndpointResponse.getData());
847+
failoverEndpointId = failoverEndpointResponse.getString("id");
848+
Assert.assertNotNull(failoverEndpointId, "Failover endpoint ID should not be null");
849+
850+
// Update the API to set the new endpoint as primary production endpoint
851+
newVersionApiDto.setPrimaryProductionEndpointId(failoverEndpointId);
852+
833853
// Remove the round-robin policy and add failover policy
834854
String failoverPolicyName = "modelFailover";
835855
Assert.assertNotNull(policyMap.get(failoverPolicyName), "Unable to find a common policy with name " + failoverPolicyName);
836-
856+
837857
Map<String, Object> failoverAttributeMap = new HashMap<>();
838858
JSONObject failoverConfigs = new JSONObject();
839-
859+
840860
// Configure production failover
841861
JSONObject productionConfig = new JSONObject();
842862
JSONObject targetModel = new JSONObject();
843863
targetModel.put("model", model1);
844864
targetModel.put("endpointId", failoverEndpointId);
845865
targetModel.put("endpointName", failoverEndpointName);
846866
productionConfig.put("targetModel", targetModel);
847-
867+
848868
List<JSONObject> fallbackModels = new ArrayList<>();
849869
JSONObject fallbackModel = new JSONObject();
850870
fallbackModel.put("model", model3);
851-
fallbackModel.put("endpointId", productionEndpointId);
852-
fallbackModel.put("endpointName", productionEndpointName);
871+
fallbackModel.put("endpointId", defaultProductionEndpointId);
872+
fallbackModel.put("endpointName", defaultProductionEndpointName);
853873
fallbackModels.add(fallbackModel);
854874
productionConfig.put("fallbackModels", fallbackModels);
855-
875+
876+
// Configure sandbox failover (empty in this case)
877+
JSONObject sandboxConfig = new JSONObject();
878+
sandboxConfig.put("targetModel", new JSONObject());
879+
sandboxConfig.put("fallbackModels", new ArrayList<JSONObject>());
880+
856881
failoverConfigs.put("production", productionConfig);
857-
failoverConfigs.put("sandbox", new JSONObject());
858-
failoverConfigs.put("requestTimeout", "10");
882+
failoverConfigs.put("sandbox", sandboxConfig);
883+
failoverConfigs.put("requestTimeout", "60");
859884
failoverConfigs.put("suspendDuration", "5");
860885

861886
String configString = failoverConfigs.toString().replace("\"", "'");
862887
failoverAttributeMap.put("failoverConfigs", configString);
863-
888+
864889
// Create failover policy
865890
List<OperationPolicyDTO> requestPolicyList = new ArrayList<>();
866891
OperationPolicyDTO failoverPolicyDTO = new OperationPolicyDTO();
@@ -869,41 +894,20 @@ public void testCreateNewVersionWithFailoverPolicy() throws Exception {
869894
failoverPolicyDTO.setPolicyId(policyMap.get(failoverPolicyName));
870895
failoverPolicyDTO.setParameters(failoverAttributeMap);
871896
requestPolicyList.add(failoverPolicyDTO);
872-
897+
873898
APIOperationPoliciesDTO apiOperationPoliciesDTO = new APIOperationPoliciesDTO();
874899
apiOperationPoliciesDTO.setRequest(requestPolicyList);
875900
newVersionApiDto.setApiPolicies(apiOperationPoliciesDTO);
876-
901+
877902
// Update the new version API with failover policy
878903
HttpResponse updateResponse = restAPIPublisher.updateAPIWithHttpInfo(newVersionApiDto);
879904
Assert.assertEquals(updateResponse.getResponseCode(), HttpStatus.SC_OK,
880-
"Failed to update new version API with failover policy");
881-
882-
// Add a new endpoint with failover URL for the new version
883-
String failoverEndpointConfig = readFile(resourcePath + "prod-endpoint-add-endpoint-config.json");
884-
JSONObject failoverEndpointConfigObj = new JSONObject(failoverEndpointConfig);
885-
JSONObject productionEndpoints = new JSONObject();
886-
productionEndpoints.put("url", endpointHost + ":" + endpointPort + failoverEndpointUrl);
887-
failoverEndpointConfigObj.put("production_endpoints", productionEndpoints);
888-
889-
Map<String, Object> failoverEndpointConfigMap = new Gson().fromJson(failoverEndpointConfigObj.toString(), Map.class);
890-
HttpResponse addFailoverEndpointResponse = restAPIPublisher.addApiEndpoint(newVersionApiId, failoverEndpointName,
891-
productionDeploymentStage, failoverEndpointConfigMap);
892-
Assert.assertEquals(addFailoverEndpointResponse.getResponseCode(), HttpStatus.SC_CREATED,
893-
"Failed to add failover endpoint to new version API");
894-
895-
// Extract failover endpoint ID from response
896-
JSONObject failoverEndpointResponse = new JSONObject(addFailoverEndpointResponse.getData());
897-
failoverEndpointId = failoverEndpointResponse.getString("id");
898-
Assert.assertNotNull(failoverEndpointId, "Failover endpoint ID should not be null");
899-
900-
// Update the API to set the new endpoint as primary production endpoint
901-
newVersionApiDto.setPrimaryProductionEndpointId(failoverEndpointId);
905+
"Failed to update new version API with failover policy and updated primary production endpoint");
902906

903907
// Create revision and deploy the new version
904908
String revisionUUID = createAPIRevisionAndDeployUsingRest(newVersionApiId, restAPIPublisher);
905909
Assert.assertNotNull(revisionUUID, "Failed to create and deploy revision for new version API");
906-
910+
907911
// Publish the new version API
908912
HttpResponse lifecycleResponse = restAPIPublisher.changeAPILifeCycleStatusToPublish(newVersionApiId, false);
909913
Assert.assertEquals(lifecycleResponse.getResponseCode(), HttpStatus.SC_OK,
@@ -914,16 +918,16 @@ public void testCreateNewVersionWithFailoverPolicy() throws Exception {
914918
newVersionApiDto.getVersion(), APIMIntegrationConstants.IS_API_EXISTS);
915919

916920
// Subscribe to the new version API
917-
SubscriptionDTO subscriptionDTO = restAPIStore.subscribeToAPI(newVersionApiId, applicationId,
918-
APIMIntegrationConstants.API_TIER.UNLIMITED);
919-
Assert.assertNotNull(subscriptionDTO, "New version API subscription failed");
921+
// SubscriptionDTO subscriptionDTO = restAPIStore.subscribeToAPI(newVersionApiId, applicationId,
922+
// APIMIntegrationConstants.API_TIER.UNLIMITED);
923+
// Assert.assertNotNull(subscriptionDTO, "New version API subscription failed");
920924

921925
// Invoke the new version API to test failover
922926
Map<String, String> requestHeaders = new HashMap<>();
923927
requestHeaders.put("ApiKey", apiKey);
924-
String invokeURL = getAPIInvocationURLHttp(newVersionContext, newVersion) + mistralAPIResource;
928+
String invokeURL = getAPIInvocationURLHttp(mistralAPIContext, newVersion) + mistralAPIResource;
925929
HttpResponse serviceResponse = HTTPSClientUtils.doPost(invokeURL, requestHeaders, mistralPayload);
926-
930+
927931
// Verify that the API call succeeded (failover should have worked)
928932
Assert.assertEquals(serviceResponse.getResponseCode(), HttpStatus.SC_OK,
929933
"API invocation should succeed due to failover");
@@ -1053,24 +1057,24 @@ private void startWiremockServer() {
10531057
.withHeader("Content-Type", "application/json")
10541058
.withBody(model1Response)));
10551059

1056-
// Stub for secured AI API (with Authorization header value of Bearer 456)
1060+
// Stub for round-robin policy: model1 request routed to model2 (80% weight)
1061+
// Using more flexible matching with containsString to handle potential JSON formatting issues
10571062
wireMockServer.stubFor(WireMock.post(urlEqualTo(mistralAPIEndpoint2 + mistralAPIResource))
10581063
.withHeader("Authorization", WireMock.matching("Bearer 456"))
1059-
// .withRequestBody(equalTo(model1))
1060-
.withRequestBody(matchingJsonPath("$.model", equalTo(model1)))
1064+
.withRequestBody(WireMock.containing(model2))
10611065
.willReturn(aResponse()
10621066
.withStatus(200)
10631067
.withHeader("Content-Type", "application/json")
1064-
.withBody(model1Response)));
1068+
.withBody(model2Response)));
10651069

1066-
// Stub for model2 (mistral-medium-latest)
1070+
// Stub for round-robin policy: model1 request routed to model3 (20% weight)
10671071
wireMockServer.stubFor(WireMock.post(urlEqualTo(mistralAPIEndpoint2 + mistralAPIResource))
10681072
.withHeader("Authorization", WireMock.matching("Bearer 456"))
1069-
.withRequestBody(matchingJsonPath("$.model", equalTo(model2)))
1073+
.withRequestBody(WireMock.containing(model3))
10701074
.willReturn(aResponse()
10711075
.withStatus(200)
10721076
.withHeader("Content-Type", "application/json")
1073-
.withBody(model2Response)));
1077+
.withBody(model3Response)));
10741078

10751079
// Stub for model3 (mistral-large-latest)
10761080
// wireMockServer.stubFor(WireMock.post(urlEqualTo(mistralAPIEndpoint2 + mistralAPIResource))
@@ -1095,9 +1099,18 @@ private void startWiremockServer() {
10951099
// .withTransformers("response-template")));
10961100

10971101
// Stub for /mistral endpoint with Bearer 123 and model3
1102+
// Using flexible matching to handle JSON formatting issues
10981103
wireMockServer.stubFor(WireMock.post(urlEqualTo(mistralAPIEndpoint + mistralAPIResource))
10991104
.withHeader("Authorization", WireMock.matching("Bearer 123"))
1100-
.withRequestBody(matchingJsonPath("$.model", equalTo(model3)))
1105+
.withRequestBody(WireMock.containing(model3))
1106+
.willReturn(aResponse()
1107+
.withStatus(200)
1108+
.withHeader("Content-Type", "application/json")
1109+
.withBody(model3Response)));
1110+
1111+
wireMockServer.stubFor(WireMock.post(urlEqualTo(mistralAPIEndpoint + '/' + newVersion + mistralAPIResource))
1112+
.withHeader("Authorization", WireMock.matching("Bearer 123"))
1113+
.withRequestBody(WireMock.containing(model3))
11011114
.willReturn(aResponse()
11021115
.withStatus(200)
11031116
.withHeader("Content-Type", "application/json")
@@ -1106,7 +1119,7 @@ private void startWiremockServer() {
11061119
// Stub for failover endpoint that returns 500 (to trigger failover)
11071120
wireMockServer.stubFor(WireMock.post(urlEqualTo(failoverEndpointUrl + mistralAPIResource))
11081121
.withHeader("Authorization", WireMock.matching("Bearer 123"))
1109-
.withRequestBody(matchingJsonPath("$.model", equalTo(model1)))
1122+
.withRequestBody(WireMock.containing(model1))
11101123
.willReturn(aResponse()
11111124
.withStatus(500)
11121125
.withHeader("Content-Type", "application/json")

0 commit comments

Comments
 (0)