Fixing MongoDB vector store support test failures

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 - 2024 the original author or authors.
+ * Copyright 2023-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,13 +18,10 @@
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.vectorstore.MongoDBAtlasVectorStore;
 import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
-import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.ObjectProvider;
-import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
-import org.springframework.boot.autoconfigure.data.mongo.MongoDataAutoConfiguration;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
 import org.springframework.core.convert.converter.Converter;
@@ -40,9 +37,10 @@
 /**
  * @author Eddú Meléndez
  * @author Christian Tzolov
+ * @author Soby Chacko
  * @since 1.0.0
  */
-@AutoConfiguration(after = MongoDataAutoConfiguration.class)
+@AutoConfiguration
 @ConditionalOnClass({ MongoDBAtlasVectorStore.class, EmbeddingModel.class, MongoTemplate.class })
 @EnableConfigurationProperties(MongoDBAtlasVectorStoreProperties.class)
 public class MongoDBAtlasVectorStoreAutoConfiguration {
@@ -92,17 +90,8 @@ public MimeType convert(String source) {
 	}
 
 	@Bean
-	public BeanPostProcessor mongoCustomConversionsPostProcessor() {
-		return new BeanPostProcessor() {
-			@Override
-			public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
-				if (bean instanceof MongoCustomConversions) {
-					return new MongoCustomConversions(
-							Arrays.asList(mimeTypeToStringConverter(), stringToMimeTypeConverter()));
-				}
-				return bean;
-			}
-		};
+	public MongoCustomConversions mongoCustomConversions() {
+		return new MongoCustomConversions(Arrays.asList(mimeTypeToStringConverter(), stringToMimeTypeConverter()));
 	}
 
 }

spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfigurationIT.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 - 2024 the original author or authors.
+ * Copyright 2023-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -53,7 +53,6 @@
  */
 @Testcontainers
 @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
-@Disabled("Disabled due to https://github.com/spring-projects/spring-ai/issues/698")
 class MongoDBAtlasVectorStoreAutoConfigurationIT {
 
 	@Container

vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStoreIT.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 - 2024 the original author or authors.
+ * Copyright 2023-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,7 +17,6 @@
 
 import com.mongodb.client.MongoClient;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.ai.document.Document;
@@ -28,10 +27,17 @@
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.context.annotation.Bean;
+import org.springframework.core.convert.converter.Converter;
 import org.springframework.data.mongodb.core.MongoTemplate;
+import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
+import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
+import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
+import org.springframework.util.MimeType;
+
 import org.testcontainers.junit.jupiter.Container;
 import org.testcontainers.junit.jupiter.Testcontainers;
 
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -42,9 +48,9 @@
 
 /**
  * @author Chris Smith
+ * @author Soby Chacko
  */
 @Testcontainers
-@Disabled("Disabled due to https://github.com/spring-projects/spring-ai/issues/698")
 class MongoDBAtlasVectorStoreIT {
 
 	@Container
@@ -201,15 +207,47 @@ public VectorStore vectorStore(MongoTemplate mongoTemplate, EmbeddingModel embed
 		}
 
 		@Bean
-		public MongoTemplate mongoTemplate(MongoClient mongoClient) {
-			return new MongoTemplate(mongoClient, "springaisample");
+		public MongoTemplate mongoTemplate(MongoClient mongoClient, MongoCustomConversions mongoCustomConversions) {
+			MongoTemplate mongoTemplate = new MongoTemplate(mongoClient, "springaisample");
+			MappingMongoConverter converter = (MappingMongoConverter) mongoTemplate.getConverter();
+			converter.setCustomConversions(mongoCustomConversions);
+			((MongoMappingContext) converter.getMappingContext())
+				.setSimpleTypeHolder(mongoCustomConversions.getSimpleTypeHolder());
+			converter.afterPropertiesSet();
+			return mongoTemplate;
 		}
 
 		@Bean
 		public EmbeddingModel embeddingModel() {
 			return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
 		}
 
+		@Bean
+		public Converter<MimeType, String> mimeTypeToStringConverter() {
+			return new Converter<MimeType, String>() {
+				@Override
+				public String convert(MimeType source) {
+					return source.toString();
+				}
+			};
+		}
+
+		@Bean
+		public Converter<String, MimeType> stringToMimeTypeConverter() {
+			return new Converter<String, MimeType>() {
+				@Override
+				public MimeType convert(String source) {
+					return MimeType.valueOf(source);
+				}
+			};
+		}
+
+		@Bean
+		public MongoCustomConversions mongoCustomConversions(Converter<MimeType, String> mimeTypeToStringConverter,
+				Converter<String, MimeType> stringToMimeTypeConverter) {
+			return new MongoCustomConversions(Arrays.asList(mimeTypeToStringConverter, stringToMimeTypeConverter));
+		}
+
 	}
 
 }

vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDbVectorStoreObservationIT.java

@@ -19,12 +19,11 @@
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
-import java.util.Vector;
 
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
 import org.springframework.ai.document.Document;
@@ -39,8 +38,14 @@
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.context.annotation.Bean;
+import org.springframework.core.convert.converter.Converter;
 import org.springframework.core.io.DefaultResourceLoader;
 import org.springframework.data.mongodb.core.MongoTemplate;
+import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
+import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
+import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
+import org.springframework.util.MimeType;
+
 import org.testcontainers.junit.jupiter.Container;
 import org.testcontainers.junit.jupiter.Testcontainers;
 
@@ -52,10 +57,10 @@
 
 /**
  * @author Christian Tzolov
+ * @author Soby Chacko
  */
 @Testcontainers
 @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
-@Disabled("Disabled due to https://github.com/spring-projects/spring-ai/issues/698")
 public class MongoDbVectorStoreObservationIT {
 
 	@Container
@@ -100,6 +105,8 @@ void observationVectorStoreAddAndQueryOperations() {
 
 			vectorStore.add(documents);
 
+			Thread.sleep(5000);
+
 			TestObservationRegistryAssert.assertThat(observationRegistry)
 				.doesNotHaveAnyRemainingCurrentObservation()
 				.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
@@ -175,15 +182,46 @@ public VectorStore vectorStore(MongoTemplate mongoTemplate, EmbeddingModel embed
 		}
 
 		@Bean
-		public MongoTemplate mongoTemplate(MongoClient mongoClient) {
-			return new MongoTemplate(mongoClient, "springaisample");
+		public MongoTemplate mongoTemplate(MongoClient mongoClient, MongoCustomConversions mongoCustomConversions) {
+			MongoTemplate mongoTemplate = new MongoTemplate(mongoClient, "springaisample");
+			MappingMongoConverter converter = (MappingMongoConverter) mongoTemplate.getConverter();
+			converter.setCustomConversions(mongoCustomConversions);
+			((MongoMappingContext) converter.getMappingContext())
+				.setSimpleTypeHolder(mongoCustomConversions.getSimpleTypeHolder());
+			converter.afterPropertiesSet();
+			return mongoTemplate;
 		}
 
 		@Bean
 		public EmbeddingModel embeddingModel() {
 			return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
 		}
 
+		@Bean
+		public Converter<MimeType, String> mimeTypeToStringConverter() {
+			return new Converter<MimeType, String>() {
+				@Override
+				public String convert(MimeType source) {
+					return source.toString();
+				}
+			};
+		}
+
+		@Bean
+		public Converter<String, MimeType> stringToMimeTypeConverter() {
+			return new Converter<String, MimeType>() {
+				@Override
+				public MimeType convert(String source) {
+					return MimeType.valueOf(source);
+				}
+			};
+		}
+
+		@Bean
+		public MongoCustomConversions mongoCustomConversions() {
+			return new MongoCustomConversions(Arrays.asList(mimeTypeToStringConverter(), stringToMimeTypeConverter()));
+		}
+
 	}
 
 }