|
16 | 16 |
|
17 | 17 | package org.springframework.ai.chat.memory.repository.jdbc; |
18 | 18 |
|
19 | | -import java.util.List; |
20 | | -import java.util.UUID; |
21 | | -import javax.sql.DataSource; |
22 | | - |
23 | | -import org.junit.jupiter.api.Test; |
24 | | - |
25 | | -import org.springframework.ai.chat.memory.ChatMemoryRepository; |
26 | | -import org.springframework.ai.chat.messages.AssistantMessage; |
27 | | -import org.springframework.ai.chat.messages.Message; |
28 | | -import org.springframework.ai.chat.messages.SystemMessage; |
29 | | -import org.springframework.ai.chat.messages.UserMessage; |
30 | | -import org.springframework.boot.SpringBootConfiguration; |
31 | 19 | import org.springframework.boot.test.context.SpringBootTest; |
32 | | -import org.springframework.context.annotation.Bean; |
33 | | -import org.springframework.jdbc.core.JdbcTemplate; |
34 | | -import org.springframework.jdbc.datasource.DataSourceTransactionManager; |
35 | 20 | import org.springframework.test.context.TestPropertySource; |
36 | 21 | import org.springframework.test.context.jdbc.Sql; |
37 | 22 |
|
38 | | -import static org.assertj.core.api.Assertions.assertThat; |
39 | | - |
40 | 23 | /** |
41 | 24 | * Integration tests for {@link JdbcChatMemoryRepository} with PostgreSQL. |
42 | 25 | * |
43 | 26 | * @author Jonathan Leijendekker |
44 | 27 | * @author Thomas Vitale |
45 | 28 | * @author Mark Pollack |
| 29 | + * @author Yanming Zhou |
46 | 30 | */ |
47 | | -@SpringBootTest(classes = JdbcChatMemoryRepositoryPostgresqlIT.TestConfiguration.class) |
| 31 | +@SpringBootTest |
48 | 32 | @TestPropertySource(properties = "spring.datasource.url=jdbc:tc:postgresql:17:///") |
49 | 33 | @Sql(scripts = "classpath:org/springframework/ai/chat/memory/repository/jdbc/schema-postgresql.sql") |
50 | 34 | class JdbcChatMemoryRepositoryPostgresqlIT extends AbstractJdbcChatMemoryRepositoryIT { |
51 | 35 |
|
52 | | - @Test |
53 | | - void repositoryWithExplicitTransactionManager() { |
54 | | - // Get the repository with explicit transaction manager |
55 | | - ChatMemoryRepository repositoryWithTxManager = TestConfiguration |
56 | | - .chatMemoryRepositoryWithTransactionManager(jdbcTemplate, jdbcTemplate.getDataSource()); |
57 | | - |
58 | | - var conversationId = UUID.randomUUID().toString(); |
59 | | - var messages = List.<Message>of(new AssistantMessage("Message with transaction manager - " + conversationId), |
60 | | - new UserMessage("User message with transaction manager - " + conversationId)); |
61 | | - |
62 | | - // Save messages using the repository with explicit transaction manager |
63 | | - repositoryWithTxManager.saveAll(conversationId, messages); |
64 | | - |
65 | | - // Verify messages were saved correctly |
66 | | - var savedMessages = repositoryWithTxManager.findByConversationId(conversationId); |
67 | | - assertThat(savedMessages).hasSize(2); |
68 | | - assertThat(savedMessages).isEqualTo(messages); |
69 | | - |
70 | | - // Verify transaction works by updating and checking atomicity |
71 | | - var newMessages = List.<Message>of(new SystemMessage("New system message - " + conversationId)); |
72 | | - repositoryWithTxManager.saveAll(conversationId, newMessages); |
73 | | - |
74 | | - // The old messages should be deleted and only the new one should exist |
75 | | - var updatedMessages = repositoryWithTxManager.findByConversationId(conversationId); |
76 | | - assertThat(updatedMessages).hasSize(1); |
77 | | - assertThat(updatedMessages).isEqualTo(newMessages); |
78 | | - } |
79 | | - |
80 | | - @SpringBootConfiguration |
81 | | - static class TestConfiguration extends BaseTestConfiguration { |
82 | | - |
83 | | - @Bean |
84 | | - ChatMemoryRepository chatMemoryRepositoryWithTxManager(JdbcTemplate jdbcTemplate, DataSource dataSource) { |
85 | | - return chatMemoryRepositoryWithTransactionManager(jdbcTemplate, dataSource); |
86 | | - } |
87 | | - |
88 | | - static ChatMemoryRepository chatMemoryRepositoryWithTransactionManager(JdbcTemplate jdbcTemplate, |
89 | | - DataSource dataSource) { |
90 | | - return JdbcChatMemoryRepository.builder() |
91 | | - .jdbcTemplate(jdbcTemplate) |
92 | | - .dialect(JdbcChatMemoryRepositoryDialect.from(dataSource)) |
93 | | - .transactionManager(new DataSourceTransactionManager(dataSource)) |
94 | | - .build(); |
95 | | - } |
96 | | - |
97 | | - } |
98 | | - |
99 | 36 | } |
0 commit comments