Skip to content

Commit 9e254a8

Browse files
jiafu1115jiafu1115
committed
fix integration test
Signed-off-by: stroller <[email protected]>
1 parent e038aa3 commit 9e254a8

File tree

3 files changed

+86
-20
lines changed

3 files changed

+86
-20
lines changed

‎auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch/src/test/java/org/springframework/ai/model/chat/memory/repository/elasticsearch/autoconfigure/ElasticSearchChatMemoryRepositoryAutoConfigurationIT.java‎

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@
1818

1919
import java.util.List;
2020
import java.util.UUID;
21+
import java.util.concurrent.TimeUnit;
2122

2223
import org.junit.jupiter.api.Test;
24+
import org.springframework.ai.chat.messages.Message;
2325
import org.testcontainers.elasticsearch.ElasticsearchContainer;
2426
import org.testcontainers.junit.jupiter.Container;
2527
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -66,25 +68,34 @@ void addAndGet() {
6668

6769
memory.saveAll(conversationId, List.of(new UserMessage("test question")));
6870

69-
assertThat(memory.findByConversationId(conversationId)).hasSize(1);
70-
assertThat(memory.findByConversationId(conversationId).get(0).getMessageType()).isEqualTo(MessageType.USER);
71-
assertThat(memory.findByConversationId(conversationId).get(0).getText()).isEqualTo("test question");
71+
sleepForSearchable();
72+
73+
List<Message> conversation = memory.findByConversationId(conversationId);
74+
assertThat(conversation).hasSize(1);
75+
assertThat(conversation.get(0).getMessageType()).isEqualTo(MessageType.USER);
76+
assertThat(conversation.get(0).getText()).isEqualTo("test question");
7277

7378
memory.deleteByConversationId(conversationId);
79+
sleepForSearchable();
7480
assertThat(memory.findByConversationId(conversationId)).isEmpty();
7581

7682
memory.saveAll(conversationId,
7783
List.of(new UserMessage("test question"), new AssistantMessage("test answer")));
78-
79-
assertThat(memory.findByConversationId(conversationId)).hasSize(2);
80-
assertThat(memory.findByConversationId(conversationId).get(0).getMessageType()).isEqualTo(MessageType.USER);
81-
assertThat(memory.findByConversationId(conversationId).get(0).getText()).isEqualTo("test question");
82-
assertThat(memory.findByConversationId(conversationId).get(1).getMessageType())
83-
.isEqualTo(MessageType.ASSISTANT);
84-
assertThat(memory.findByConversationId(conversationId).get(1).getText()).isEqualTo("test answer");
84+
sleepForSearchable();
85+
86+
conversation = memory.findByConversationId(conversationId);
87+
assertThat(conversation).hasSize(2);
88+
assertThat(conversation.get(0).getMessageType()).isEqualTo(MessageType.USER);
89+
assertThat(conversation.get(0).getText()).isEqualTo("test question");
90+
assertThat(conversation.get(1).getMessageType()).isEqualTo(MessageType.ASSISTANT);
91+
assertThat(conversation.get(1).getText()).isEqualTo("test answer");
8592
});
8693
}
8794

95+
private static void sleepForSearchable() throws InterruptedException {
96+
TimeUnit.SECONDS.sleep(2);
97+
}
98+
8899
@Test
89100
void propertiesConfiguration() {
90101
this.contextRunner
@@ -107,6 +118,8 @@ void findConversationIds() {
107118
memory.saveAll(conversationId1, List.of(new UserMessage("test question 1")));
108119
memory.saveAll(conversationId2, List.of(new UserMessage("test question 2")));
109120

121+
sleepForSearchable();
122+
110123
List<String> conversationIds = memory.findConversationIds();
111124
assertThat(conversationIds).contains(conversationId1, conversationId2);
112125
});

‎memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/main/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepository.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import java.util.stream.Collectors;
2626

2727
import co.elastic.clients.elasticsearch.ElasticsearchClient;
28+
import co.elastic.clients.elasticsearch._types.SortOrder;
2829
import co.elastic.clients.elasticsearch.core.BulkRequest;
2930
import co.elastic.clients.elasticsearch.core.SearchResponse;
3031
import co.elastic.clients.elasticsearch.core.search.Hit;
@@ -100,8 +101,7 @@ public List<Message> findByConversationId(String conversationId) {
100101
try {
101102
SearchResponse<Map> response = this.client.search(s -> s.index(this.indexName)
102103
.query(q -> q.term(t -> t.field("conversationId").value(conversationId)))
103-
.sort(sort -> sort
104-
.field(f -> f.field("sequenceNumber").order(co.elastic.clients.elasticsearch._types.SortOrder.Asc)))
104+
.sort(sort -> sort.field(f -> f.field("sequenceNumber").order(SortOrder.Asc)))
105105
.size(10000), Map.class);
106106

107107
return response.hits()

‎memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/test/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepositoryIT.java‎

Lines changed: 61 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,10 @@ void add_shouldInsertSingleMessage(String content, MessageType messageType) {
9191
default -> throw new IllegalArgumentException("Type not supported: " + messageType);
9292
};
9393

94+
java.time.Instant beforeSave = java.time.Instant.now();
9495
chatMemory.saveAll(conversationId, List.of(message));
96+
java.time.Instant afterSave = java.time.Instant.now();
97+
9598
sleepForSearchable();
9699
assertThat(chatMemory.findConversationIds()).isNotEmpty();
97100
assertThat(chatMemory.findConversationIds()).contains(conversationId);
@@ -100,6 +103,15 @@ void add_shouldInsertSingleMessage(String content, MessageType messageType) {
100103
assertThat(retrievedMessages).hasSize(1);
101104
assertThat(retrievedMessages.get(0).getText()).isEqualTo(content);
102105
assertThat(retrievedMessages.get(0).getMessageType()).isEqualTo(messageType);
106+
107+
// Verify timestamp was automatically added
108+
Object timestampObj = retrievedMessages.get(0)
109+
.getMetadata()
110+
.get(ElasticSearchChatMemoryRepository.CONVERSATION_TS);
111+
assertThat(timestampObj).as("Timestamp should be automatically added").isNotNull();
112+
assertThat(timestampObj).isInstanceOf(java.time.Instant.class);
113+
java.time.Instant timestamp = (java.time.Instant) timestampObj;
114+
assertThat(timestamp).isBetween(beforeSave.minusSeconds(1), afterSave.plusSeconds(1));
103115
});
104116
}
105117

@@ -115,17 +127,21 @@ void shouldSaveAndRetrieveMultipleMessages() {
115127

116128
var conversationId = UUID.randomUUID().toString();
117129

130+
// Use 5 messages to verify sequenceNumber preserves order
118131
List<Message> messages = List.of(new SystemMessage("System message"), new UserMessage("User message"),
119-
new AssistantMessage("Assistant message"));
132+
new AssistantMessage("Assistant message"), new UserMessage("Second user message"),
133+
new AssistantMessage("Second assistant message"));
120134

135+
java.time.Instant beforeSave = java.time.Instant.now();
121136
chatMemory.saveAll(conversationId, messages);
137+
java.time.Instant afterSave = java.time.Instant.now();
122138

123139
sleepForSearchable();
124140

125141
List<Message> retrievedMessages = chatMemory.findByConversationId(conversationId);
126-
assertThat(retrievedMessages).hasSize(3);
142+
assertThat(retrievedMessages).hasSize(5);
127143

128-
// Messages should be in the same order they were saved
144+
// Verify sequenceNumber preserves message order
129145
assertThat(retrievedMessages.get(0).getText()).isEqualTo("System message");
130146
assertThat(retrievedMessages.get(0).getMessageType()).isEqualTo(MessageType.SYSTEM);
131147

@@ -134,6 +150,31 @@ void shouldSaveAndRetrieveMultipleMessages() {
134150

135151
assertThat(retrievedMessages.get(2).getText()).isEqualTo("Assistant message");
136152
assertThat(retrievedMessages.get(2).getMessageType()).isEqualTo(MessageType.ASSISTANT);
153+
154+
assertThat(retrievedMessages.get(3).getText()).isEqualTo("Second user message");
155+
assertThat(retrievedMessages.get(3).getMessageType()).isEqualTo(MessageType.USER);
156+
157+
assertThat(retrievedMessages.get(4).getText()).isEqualTo("Second assistant message");
158+
assertThat(retrievedMessages.get(4).getMessageType()).isEqualTo(MessageType.ASSISTANT);
159+
160+
// Verify timestamp consistency: all messages saved together should have the
161+
// same timestamp
162+
java.time.Instant firstTimestamp = null;
163+
for (Message message : retrievedMessages) {
164+
Object timestampObj = message.getMetadata().get(ElasticSearchChatMemoryRepository.CONVERSATION_TS);
165+
assertThat(timestampObj).isNotNull();
166+
167+
java.time.Instant timestamp = (java.time.Instant) timestampObj;
168+
assertThat(timestamp).isBetween(beforeSave.minusSeconds(1), afterSave.plusSeconds(1));
169+
170+
if (firstTimestamp == null) {
171+
firstTimestamp = timestamp;
172+
}
173+
else {
174+
// All messages in same saveAll should have identical timestamp
175+
assertThat(timestamp).isEqualTo(firstTimestamp);
176+
}
177+
}
137178
});
138179
}
139180

@@ -143,18 +184,24 @@ void shouldReplaceExistingMessages() {
143184
var chatMemory = context.getBean(ChatMemoryRepository.class);
144185
var conversationId = UUID.randomUUID().toString();
145186

146-
// Save initial messages
187+
// Save initial messages (3 messages)
147188
List<Message> initialMessages = List.of(new UserMessage("Initial user message"),
148-
new AssistantMessage("Initial assistant message"));
189+
new AssistantMessage("Initial assistant message"), new UserMessage("Initial third message"));
149190
chatMemory.saveAll(conversationId, initialMessages);
150191

151192
sleepForSearchable();
152193

153-
// Verify initial save
194+
// Verify initial save and capture timestamp
154195
List<Message> retrievedMessages = chatMemory.findByConversationId(conversationId);
155-
assertThat(retrievedMessages).hasSize(2);
196+
assertThat(retrievedMessages).hasSize(3);
197+
java.time.Instant initialTimestamp = (java.time.Instant) retrievedMessages.get(0)
198+
.getMetadata()
199+
.get(ElasticSearchChatMemoryRepository.CONVERSATION_TS);
156200

157-
// Replace with new messages
201+
// Wait to ensure timestamp difference
202+
TimeUnit.MILLISECONDS.sleep(100);
203+
204+
// Replace with new messages (2 messages - verify sequenceNumber restarts)
158205
List<Message> newMessages = List.of(new SystemMessage("New system message"),
159206
new UserMessage("New user message"));
160207
chatMemory.saveAll(conversationId, newMessages);
@@ -166,6 +213,12 @@ void shouldReplaceExistingMessages() {
166213
assertThat(retrievedMessages).hasSize(2);
167214
assertThat(retrievedMessages.get(0).getText()).isEqualTo("New system message");
168215
assertThat(retrievedMessages.get(1).getText()).isEqualTo("New user message");
216+
217+
// Verify timestamp was updated (should be after initial timestamp)
218+
java.time.Instant updatedTimestamp = (java.time.Instant) retrievedMessages.get(0)
219+
.getMetadata()
220+
.get(ElasticSearchChatMemoryRepository.CONVERSATION_TS);
221+
assertThat(updatedTimestamp).isAfter(initialTimestamp);
169222
});
170223
}
171224

0 commit comments

Comments
 (0)