Yevhenii
9 months ago
5 changed files with 228 additions and 44 deletions
@ -0,0 +1,197 @@ |
|||
/** |
|||
* Copyright © 2016-2025 The Thingsboard Authors |
|||
* |
|||
* Licensed under the Apache License, Version 2.0 (the "License"); |
|||
* you may not use this file except in compliance with the License. |
|||
* You may obtain a copy of the License at |
|||
* |
|||
* http://www.apache.org/licenses/LICENSE-2.0
|
|||
* |
|||
* Unless required by applicable law or agreed to in writing, software |
|||
* distributed under the License is distributed on an "AS IS" BASIS, |
|||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|||
* See the License for the specific language governing permissions and |
|||
* limitations under the License. |
|||
*/ |
|||
package org.thingsboard.server.edge; |
|||
|
|||
import com.datastax.oss.driver.api.core.uuid.Uuids; |
|||
import com.google.protobuf.AbstractMessage; |
|||
import com.google.protobuf.InvalidProtocolBufferException; |
|||
import org.junit.Assert; |
|||
import org.junit.Test; |
|||
import org.thingsboard.common.util.JacksonUtil; |
|||
import org.thingsboard.server.common.data.ai.AiModel; |
|||
import org.thingsboard.server.common.data.ai.model.chat.OpenAiChatModelConfig; |
|||
import org.thingsboard.server.common.data.ai.provider.OpenAiProviderConfig; |
|||
import org.thingsboard.server.common.data.id.TenantId; |
|||
import org.thingsboard.server.dao.service.DaoSqlTest; |
|||
import org.thingsboard.server.gen.edge.v1.AiModelUpdateMsg; |
|||
import org.thingsboard.server.gen.edge.v1.UpdateMsgType; |
|||
import org.thingsboard.server.gen.edge.v1.UplinkMsg; |
|||
import org.thingsboard.server.gen.edge.v1.UplinkResponseMsg; |
|||
|
|||
import java.util.Optional; |
|||
import java.util.UUID; |
|||
|
|||
import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; |
|||
|
|||
@DaoSqlTest |
|||
public class AiModelEdgeTest extends AbstractEdgeTest { |
|||
|
|||
private static final String DEFAULT_AI_MODEL_NAME = "Edge Test AiModel"; |
|||
private static final String UPDATED_AI_MODEL_NAME = "Updated Edge Test AiModel"; |
|||
|
|||
@Test |
|||
public void testAiModel_create_update_delete() throws Exception { |
|||
// create AiModel
|
|||
AiModel aiModel = createSimpleAiModel(DEFAULT_AI_MODEL_NAME); |
|||
|
|||
edgeImitator.expectMessageAmount(1); |
|||
AiModel savedAiModel = doPost("/api/ai/model", aiModel, AiModel.class); |
|||
Assert.assertTrue(edgeImitator.waitForMessages()); |
|||
|
|||
AbstractMessage latestMessage = edgeImitator.getLatestMessage(); |
|||
Assert.assertTrue(latestMessage instanceof AiModelUpdateMsg); |
|||
AiModelUpdateMsg aiModelUpdateMsg = (AiModelUpdateMsg) latestMessage; |
|||
Assert.assertEquals(UpdateMsgType.ENTITY_CREATED_RPC_MESSAGE, aiModelUpdateMsg.getMsgType()); |
|||
Assert.assertEquals(savedAiModel.getUuidId().getMostSignificantBits(), aiModelUpdateMsg.getIdMSB()); |
|||
Assert.assertEquals(savedAiModel.getUuidId().getLeastSignificantBits(), aiModelUpdateMsg.getIdLSB()); |
|||
AiModel aiModelFromMsg = JacksonUtil.fromString(aiModelUpdateMsg.getEntity(), AiModel.class, true); |
|||
Assert.assertNotNull(aiModelFromMsg); |
|||
|
|||
Assert.assertEquals(DEFAULT_AI_MODEL_NAME, aiModelFromMsg.getName()); |
|||
Assert.assertEquals(savedAiModel.getTenantId(), aiModelFromMsg.getTenantId()); |
|||
|
|||
// update AiModel
|
|||
edgeImitator.expectMessageAmount(1); |
|||
savedAiModel.setName(UPDATED_AI_MODEL_NAME); |
|||
savedAiModel = doPost("/api/ai/model", savedAiModel, AiModel.class); |
|||
Assert.assertTrue(edgeImitator.waitForMessages()); |
|||
|
|||
latestMessage = edgeImitator.getLatestMessage(); |
|||
Assert.assertTrue(latestMessage instanceof AiModelUpdateMsg); |
|||
aiModelUpdateMsg = (AiModelUpdateMsg) latestMessage; |
|||
aiModelFromMsg = JacksonUtil.fromString(aiModelUpdateMsg.getEntity(), AiModel.class, true); |
|||
Assert.assertNotNull(aiModelFromMsg); |
|||
Assert.assertEquals(UpdateMsgType.ENTITY_UPDATED_RPC_MESSAGE, aiModelUpdateMsg.getMsgType()); |
|||
Assert.assertEquals(UPDATED_AI_MODEL_NAME, aiModelFromMsg.getName()); |
|||
|
|||
// delete AiModel
|
|||
edgeImitator.expectMessageAmount(1); |
|||
doDelete("/api/ai/model/" + savedAiModel.getUuidId()) |
|||
.andExpect(status().isOk()); |
|||
Assert.assertTrue(edgeImitator.waitForMessages()); |
|||
|
|||
latestMessage = edgeImitator.getLatestMessage(); |
|||
Assert.assertTrue(latestMessage instanceof AiModelUpdateMsg); |
|||
aiModelUpdateMsg = (AiModelUpdateMsg) latestMessage; |
|||
Assert.assertEquals(UpdateMsgType.ENTITY_DELETED_RPC_MESSAGE, aiModelUpdateMsg.getMsgType()); |
|||
Assert.assertEquals(savedAiModel.getUuidId().getMostSignificantBits(), aiModelUpdateMsg.getIdMSB()); |
|||
Assert.assertEquals(savedAiModel.getUuidId().getLeastSignificantBits(), aiModelUpdateMsg.getIdLSB()); |
|||
} |
|||
|
|||
@Test |
|||
public void testSendAiModelToCloud() throws Exception { |
|||
AiModel aiModel = createSimpleAiModel(DEFAULT_AI_MODEL_NAME); |
|||
UUID uuid = Uuids.timeBased(); |
|||
UplinkMsg uplinkMsg = getUplinkMsg(uuid, aiModel, UpdateMsgType.ENTITY_CREATED_RPC_MESSAGE); |
|||
|
|||
checkAiModelOnCloud(uplinkMsg, uuid, aiModel.getName()); |
|||
} |
|||
|
|||
@Test |
|||
public void testUpdateAiModelNameOnCloud() throws Exception { |
|||
AiModel aiModel = createSimpleAiModel(DEFAULT_AI_MODEL_NAME); |
|||
UUID uuid = Uuids.timeBased(); |
|||
UplinkMsg uplinkMsg = getUplinkMsg(uuid, aiModel, UpdateMsgType.ENTITY_CREATED_RPC_MESSAGE); |
|||
|
|||
checkAiModelOnCloud(uplinkMsg, uuid, aiModel.getName()); |
|||
|
|||
aiModel.setName(UPDATED_AI_MODEL_NAME); |
|||
UplinkMsg updatedUplinkMsg = getUplinkMsg(uuid, aiModel, UpdateMsgType.ENTITY_UPDATED_RPC_MESSAGE); |
|||
|
|||
checkAiModelOnCloud(updatedUplinkMsg, uuid, aiModel.getName()); |
|||
} |
|||
|
|||
@Test |
|||
public void testAiModelToCloudWithNameThatAlreadyExistsOnCloud() throws Exception { |
|||
AiModel aiModel = createSimpleAiModel(DEFAULT_AI_MODEL_NAME); |
|||
|
|||
edgeImitator.expectMessageAmount(1); |
|||
AiModel savedAiModel = doPost("/api/ai/model", aiModel, AiModel.class); |
|||
Assert.assertTrue(edgeImitator.waitForMessages()); |
|||
|
|||
UUID uuid = Uuids.timeBased(); |
|||
|
|||
UplinkMsg uplinkMsg = getUplinkMsg(uuid, aiModel, UpdateMsgType.ENTITY_CREATED_RPC_MESSAGE); |
|||
|
|||
edgeImitator.expectResponsesAmount(1); |
|||
edgeImitator.expectMessageAmount(1); |
|||
edgeImitator.sendUplinkMsg(uplinkMsg); |
|||
|
|||
Assert.assertTrue(edgeImitator.waitForResponses()); |
|||
Assert.assertTrue(edgeImitator.waitForMessages()); |
|||
|
|||
Optional<AiModelUpdateMsg> aiModelUpdateMsgOpt = edgeImitator.findMessageByType(AiModelUpdateMsg.class); |
|||
Assert.assertTrue(aiModelUpdateMsgOpt.isPresent()); |
|||
AiModelUpdateMsg latestAiModelUpdateMsg = aiModelUpdateMsgOpt.get(); |
|||
AiModel aiModelFromMsg = JacksonUtil.fromString(latestAiModelUpdateMsg.getEntity(), AiModel.class, true); |
|||
Assert.assertNotNull(aiModelFromMsg); |
|||
Assert.assertNotEquals(DEFAULT_AI_MODEL_NAME, aiModelFromMsg.getName()); |
|||
|
|||
Assert.assertNotEquals(savedAiModel.getUuidId(), uuid); |
|||
|
|||
AiModel aiModelFromCloud = doGet("/api/ai/model/" + uuid, AiModel.class); |
|||
Assert.assertNotNull(aiModelFromCloud); |
|||
Assert.assertNotEquals(DEFAULT_AI_MODEL_NAME, aiModelFromCloud.getName()); |
|||
} |
|||
|
|||
private AiModel createSimpleAiModel(String name) { |
|||
AiModel aiModel = new AiModel(); |
|||
aiModel.setTenantId(tenantId); |
|||
aiModel.setName(name); |
|||
aiModel.setConfiguration(OpenAiChatModelConfig.builder() |
|||
.providerConfig(new OpenAiProviderConfig("test-api-key")) |
|||
.modelId("gpt-4o") |
|||
.temperature(0.5) |
|||
.topP(0.3) |
|||
.frequencyPenalty(0.1) |
|||
.presencePenalty(0.2) |
|||
.maxOutputTokens(1000) |
|||
.timeoutSeconds(60) |
|||
.maxRetries(2) |
|||
.build()); |
|||
return aiModel; |
|||
} |
|||
|
|||
private UplinkMsg getUplinkMsg(UUID uuid, AiModel aiModel, UpdateMsgType updateMsgType) throws InvalidProtocolBufferException { |
|||
UplinkMsg.Builder uplinkMsgBuilder = UplinkMsg.newBuilder(); |
|||
AiModelUpdateMsg.Builder aiModelUpdateMsgBuilder = AiModelUpdateMsg.newBuilder(); |
|||
aiModelUpdateMsgBuilder.setIdMSB(uuid.getMostSignificantBits()); |
|||
aiModelUpdateMsgBuilder.setIdLSB(uuid.getLeastSignificantBits()); |
|||
aiModelUpdateMsgBuilder.setEntity(JacksonUtil.toString(aiModel)); |
|||
aiModelUpdateMsgBuilder.setMsgType(updateMsgType); |
|||
testAutoGeneratedCodeByProtobuf(aiModelUpdateMsgBuilder); |
|||
uplinkMsgBuilder.addAiModelUpdateMsg(aiModelUpdateMsgBuilder.build()); |
|||
|
|||
testAutoGeneratedCodeByProtobuf(uplinkMsgBuilder); |
|||
|
|||
return uplinkMsgBuilder.build(); |
|||
} |
|||
|
|||
private void checkAiModelOnCloud(UplinkMsg uplinkMsg, UUID uuid, String resourceTitle) throws Exception { |
|||
edgeImitator.expectResponsesAmount(1); |
|||
edgeImitator.sendUplinkMsg(uplinkMsg); |
|||
|
|||
Assert.assertTrue(edgeImitator.waitForResponses()); |
|||
|
|||
UplinkResponseMsg latestResponseMsg = edgeImitator.getLatestResponseMsg(); |
|||
Assert.assertTrue(latestResponseMsg.getSuccess()); |
|||
|
|||
AiModel aiModel = doGet("/api/ai/model/" + uuid, AiModel.class); |
|||
Assert.assertNotNull(aiModel); |
|||
Assert.assertEquals(resourceTitle, aiModel.getName()); |
|||
} |
|||
|
|||
} |
|||
Loading…
Reference in new issue