Instrument MongoCollection.countDocuments for better tracking

client-java/instrumentation/src/main/java/org/evomaster/client/java/instrumentation/coverage/methodreplacement/thirdpartyclasses/MongoCollectionClassReplacement.java

@@ -51,6 +51,53 @@ public static <TResult> Object find(Object mongoCollection, @ThirdPartyCast(actu
         return handleFind("findClientSessionBsonResultClass", mongoCollection, Arrays.asList(clientSession, filter, resultClass), filter);
     }
 
+    @Replacement(replacingStatic = false, type = ReplacementType.TRACKER, id = "countDocumentsBson", usageFilter = UsageFilter.ANY, category = ReplacementCategory.MONGO)
+    public static long countDocuments(Object mongoCollection, @ThirdPartyCast(actualType = "org.bson.conversions.Bson") Object filter) {
+        return handleCount("countDocumentsBson", mongoCollection, Arrays.asList(filter), filter);
+    }
+
+    @Replacement(replacingStatic = false, type = ReplacementType.TRACKER, id = "countDocumentsBsonCountOptions", usageFilter = UsageFilter.ANY, category = ReplacementCategory.MONGO)
+    public static long countDocuments(Object mongoCollection, @ThirdPartyCast(actualType = "org.bson.conversions.Bson") Object filter, @ThirdPartyCast(actualType = "com.mongodb.client.model.CountOptions") Object options) {
+        return handleCount("countDocumentsBsonCountOptions", mongoCollection, Arrays.asList(filter, options), filter);
+    }
+
+    @Replacement(replacingStatic = false, type = ReplacementType.TRACKER, id = "countDocumentsClientSessionBson", usageFilter = UsageFilter.ANY, category = ReplacementCategory.MONGO)
+    public static long countDocuments_EM_0(Object mongoCollection, @ThirdPartyCast(actualType = "com.mongodb.client.ClientSession") Object clientSession, @ThirdPartyCast(actualType = "org.bson.conversions.Bson") Object filter) {
+        return handleCount("countDocumentsClientSessionBson", mongoCollection, Arrays.asList(clientSession, filter), filter);
+    }
+
+    @Replacement(replacingStatic = false, type = ReplacementType.TRACKER, id = "countDocumentsClientSessionBsonCountOptions", usageFilter = UsageFilter.ANY, category = ReplacementCategory.MONGO)
+    public static long countDocuments(Object mongoCollection, @ThirdPartyCast(actualType = "com.mongodb.client.ClientSession") Object clientSession, @ThirdPartyCast(actualType = "org.bson.conversions.Bson") Object filter, @ThirdPartyCast(actualType = "com.mongodb.client.model.CountOptions") Object options) {
+        return handleCount("countDocumentsClientSessionBsonCountOptions", mongoCollection, Arrays.asList(clientSession, filter, options), filter);
+    }
+
+    /**
+     * Unlike find, count is executed eagerly, so we know right away whether the filter was valid
+     */
+    private static long handleCount(String id, Object mongoCollection, List<Object> args, Object filter) {
+        long start = System.currentTimeMillis();
+        Method countMethod = retrieveFindMethod(id, mongoCollection);
+        try {
+            long result = (long) countMethod.invoke(mongoCollection, args.toArray());
+            long end = System.currentTimeMillis();
+            handleMongo(mongoCollection, filter, true, end  - start);
+            return result;
+        } catch (IllegalAccessException e) {
+            throw new RuntimeException(e);
+        } catch (InvocationTargetException e) {
+            long end = System.currentTimeMillis();
+            handleMongo(mongoCollection, filter, false, end - start);
+            Throwable cause = e.getCause();
+            if (cause instanceof RuntimeException) {
+                throw (RuntimeException) cause;
+            }
+            if (cause instanceof Error) {
+                throw (Error) cause;
+            }
+            throw new RuntimeException(cause);
+        }
+    }
+
     private static Object handleFind(String id, Object mongoCollection, List<Object> args, Object query) {
         long start = System.currentTimeMillis();
         try {

client-java/instrumentation/src/test/java/org/evomaster/client/java/instrumentation/coverage/methodreplacement/thirdpartyclasses/MongoCollectionClassReplacementTest.java

@@ -349,5 +349,105 @@ public void testFindWithInvalidFilter() {
 
     }
 
+    private BsonDocument assertSingleRecordedCommand(boolean expectedSuccess) {
+        List<AdditionalInfo> additionalInfoList = ExecutionTracer.exposeAdditionalInfoList();
+        assertEquals(1, additionalInfoList.size());
+        Set<MongoFindCommand> mongoFindCommands = additionalInfoList.get(0).getMongoInfoData();
+        assertEquals(1, mongoFindCommands.size());
+
+        MongoFindCommand mongoFindCommand = mongoFindCommands.iterator().next();
+        assertEquals(COLLECTION_NAME, mongoFindCommand.getCollectionName());
+        assertEquals(DATABASE_NAME, mongoFindCommand.getDatabaseName());
+        assertEquals(expectedSuccess, mongoFindCommand.isSuccessfullyExecuted());
+        assertNotNull(mongoFindCommand.getQuery());
+        return (BsonDocument) mongoFindCommand.getQuery();
+    }
+
+    private void insertPerson(String name, int age) {
+        getMongoCollection().insertOne(new Document("name", name).append("age", age));
+    }
+
+    @Test
+    public void testCountDocumentsWithFilter() {
+        final MongoCollection<Document> collection = getMongoCollection();
+        insertPerson("John Doe", 30);
+        insertPerson("Jane Doe", 25);
+
+        ExecutionTracer.setExecutingInitMongo(false);
+        long count = MongoCollectionClassReplacement.countDocuments(collection, new Document("age", 30));
+
+        assertEquals(1, count);
+        BsonDocument retrievedQuery = assertSingleRecordedCommand(true);
+        assertEquals(30, retrievedQuery.getInt32("age").getValue());
+    }
+
+    @Test
+    public void testCountDocumentsWithFilterNoMatch() {
+        final MongoCollection<Document> collection = getMongoCollection();
+
+        ExecutionTracer.setExecutingInitMongo(false);
+        long count = MongoCollectionClassReplacement.countDocuments(collection, new Document("age", 99));
+
+        assertEquals(0, count);
+        BsonDocument retrievedQuery = assertSingleRecordedCommand(true);
+        assertEquals(99, retrievedQuery.getInt32("age").getValue());
+    }
+
+    @Test
+    public void testCountDocumentsWithFilterAndOptions() {
+        final MongoCollection<Document> collection = getMongoCollection();
+        insertPerson("John Doe", 30);
+        insertPerson("Jane Doe", 30);
+
+        ExecutionTracer.setExecutingInitMongo(false);
+        long count = MongoCollectionClassReplacement.countDocuments(collection, new Document("age", 30), new com.mongodb.client.model.CountOptions().limit(1));
 
+        assertEquals(1, count);
+        BsonDocument retrievedQuery = assertSingleRecordedCommand(true);
+        assertEquals(30, retrievedQuery.getInt32("age").getValue());
+    }
+
+    @Test
+    public void testCountDocumentsWithClientSessionAndFilter() {
+        try (ClientSession clientSession = mongoClient.startSession()) {
+            final MongoCollection<Document> collection = getMongoCollection();
+            insertPerson("John Doe", 23);
+
+            ExecutionTracer.setExecutingInitMongo(false);
+            long count = MongoCollectionClassReplacement.countDocuments_EM_0(collection, clientSession, new Document("age", 23));
+
+            assertEquals(1, count);
+            BsonDocument retrievedQuery = assertSingleRecordedCommand(true);
+            assertEquals(23, retrievedQuery.getInt32("age").getValue());
+        }
+    }
+
+    @Test
+    public void testCountDocumentsWithClientSessionFilterAndOptions() {
+        try (ClientSession clientSession = mongoClient.startSession()) {
+            final MongoCollection<Document> collection = getMongoCollection();
+            insertPerson("John Doe", 23);
+
+            ExecutionTracer.setExecutingInitMongo(false);
+            long count = MongoCollectionClassReplacement.countDocuments(collection, clientSession, new Document("age", 23), new com.mongodb.client.model.CountOptions().limit(1));
+
+            assertEquals(1, count);
+            BsonDocument retrievedQuery = assertSingleRecordedCommand(true);
+            assertEquals(23, retrievedQuery.getInt32("age").getValue());
+        }
+    }
+
+    @Test
+    public void testCountDocumentsWithInvalidFilter() {
+        final MongoCollection<Document> collection = getMongoCollection();
+        insertPerson("John Doe", 30);
+
+        Document invalidFilter = new Document("tags", new Document("$size", -1));
+
+        ExecutionTracer.setExecutingInitMongo(false);
+        assertThrows(com.mongodb.MongoException.class, () -> MongoCollectionClassReplacement.countDocuments(collection, invalidFilter));
+
+        BsonDocument retrievedQuery = assertSingleRecordedCommand(false);
+        assertEquals(-1, retrievedQuery.getDocument("tags").getInt32("$size").getValue());
+    }
 }

core-tests/e2e-tests/spring/spring-rest-mongo/src/main/java/com/mongo/count/MongoCountApp.java

@@ -0,0 +1,20 @@
+package com.mongo.count;
+
+import com.mongo.SwaggerConfiguration;
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.autoconfigure.security.servlet.SecurityAutoConfiguration;
+import springfox.documentation.swagger2.annotations.EnableSwagger2;
+
+@EnableSwagger2
+@SpringBootApplication(exclude = SecurityAutoConfiguration.class)
+public class MongoCountApp extends SwaggerConfiguration {
+    public MongoCountApp() {
+        super("mongocount");
+    }
+
+    public static void main(String[] args) {
+        SpringApplication.run(MongoCountApp.class, args);
+    }
+
+}

core-tests/e2e-tests/spring/spring-rest-mongo/src/main/java/com/mongo/count/MongoCountController.java

@@ -0,0 +1,112 @@
+package com.mongo.count;
+
+import com.mongodb.client.ClientSession;
+import com.mongodb.client.MongoClient;
+import com.mongodb.client.MongoCollection;
+import com.mongodb.client.model.CountOptions;
+import com.mongodb.client.model.Filters;
+import org.bson.conversions.Bson;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.data.mongodb.core.MongoTemplate;
+import org.springframework.http.ResponseEntity;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RequestParam;
+import org.springframework.web.bind.annotation.RestController;
+
+/**
+ * Counts documents directly through the Mongo driver (MongoCollection.countDocuments),
+ * not through Spring Data, so the filter is only visible at the driver level.
+ * There is one endpoint for each overload of countDocuments.
+ */
+@RestController
+@RequestMapping(path = "/mongocount")
+public class MongoCountController {
+
+    /*
+        Each endpoint counts on its own collection, so that a document that
+        makes one endpoint return 200 cannot accidentally do it for another one.
+     */
+    private static final String COUNT_ALL = "countAllCollection";
+    private static final String COUNT_BSON = "countBsonCollection";
+    private static final String COUNT_BSON_OPTIONS = "countBsonOptionsCollection";
+    private static final String COUNT_SESSION = "countSessionCollection";
+    private static final String COUNT_SESSION_BSON = "countSessionBsonCollection";
+    private static final String COUNT_SESSION_BSON_OPTIONS = "countSessionBsonOptionsCollection";
+
+    @Autowired
+    private MongoTemplate mongoTemplate;
+
+    @Autowired
+    private MongoClient mongoClient;
+
+    private MongoCollection<MongoCountData> collection(String collectionName) {
+        return mongoTemplate.getDb().getCollection(collectionName, MongoCountData.class);
+    }
+
+    private static Bson filter(String city, int minAge) {
+        return Filters.and(Filters.eq("city", city), Filters.gte("age", minAge));
+    }
+
+    private static ResponseEntity<Void> response(long count) {
+        return ResponseEntity.status(count > 0 ? 200 : 404).build();
+    }
+
+    /**
+     * countDocuments()
+     */
+    @GetMapping("countAll")
+    public ResponseEntity<Void> countAll() {
+        return response(collection(COUNT_ALL).countDocuments());
+    }
+
+    /**
+     * countDocuments(Bson)
+     */
+    @GetMapping("countByCityAndAge")
+    public ResponseEntity<Void> countByCityAndAge(@RequestParam(name = "city") String city,
+                                                  @RequestParam(name = "minAge") int minAge) {
+        return response(collection(COUNT_BSON).countDocuments(filter(city, minAge)));
+    }
+
+    /**
+     * countDocuments(Bson, CountOptions)
+     */
+    @GetMapping("countByCityAndAgeWithOptions")
+    public ResponseEntity<Void> countByCityAndAgeWithOptions(@RequestParam(name = "city") String city,
+                                                             @RequestParam(name = "minAge") int minAge) {
+        return response(collection(COUNT_BSON_OPTIONS).countDocuments(filter(city, minAge), new CountOptions().limit(1)));
+    }
+
+    /**
+     * countDocuments(ClientSession)
+     */
+    @GetMapping("countAllInSession")
+    public ResponseEntity<Void> countAllInSession() {
+        try (ClientSession session = mongoClient.startSession()) {
+            return response(collection(COUNT_SESSION).countDocuments(session));
+        }
+    }
+
+    /**
+     * countDocuments(ClientSession, Bson)
+     */
+    @GetMapping("countByCityAndAgeInSession")
+    public ResponseEntity<Void> countByCityAndAgeInSession(@RequestParam(name = "city") String city,
+                                                           @RequestParam(name = "minAge") int minAge) {
+        try (ClientSession session = mongoClient.startSession()) {
+            return response(collection(COUNT_SESSION_BSON).countDocuments(session, filter(city, minAge)));
+        }
+    }
+
+    /**
+     * countDocuments(ClientSession, Bson, CountOptions)
+     */
+    @GetMapping("countByCityAndAgeInSessionWithOptions")
+    public ResponseEntity<Void> countByCityAndAgeInSessionWithOptions(@RequestParam(name = "city") String city,
+                                                                      @RequestParam(name = "minAge") int minAge) {
+        try (ClientSession session = mongoClient.startSession()) {
+            return response(collection(COUNT_SESSION_BSON_OPTIONS).countDocuments(session, filter(city, minAge), new CountOptions().limit(1)));
+        }
+    }
+}

core-tests/e2e-tests/spring/spring-rest-mongo/src/main/java/com/mongo/count/MongoCountData.java

@@ -0,0 +1,21 @@
+package com.mongo.count;
+
+import org.springframework.data.annotation.Id;
+
+public class MongoCountData {
+
+    @Id
+    public String id;
+
+    public String city;
+
+    public int age;
+
+    public MongoCountData() {
+    }
+
+    public MongoCountData(String city, int age) {
+        this.city = city;
+        this.age = age;
+    }
+}

core-tests/e2e-tests/spring/spring-rest-mongo/src/test/java/com/foo/spring/rest/mongo/count/MongoCountAppController.java

@@ -0,0 +1,15 @@
+package com.foo.spring.rest.mongo.count;
+
+import com.foo.spring.rest.mongo.MongoController;
+import com.mongo.count.MongoCountApp;
+
+public class MongoCountAppController extends MongoController {
+    public MongoCountAppController() {
+        super("count", MongoCountApp.class);
+    }
+
+    @Override
+    public String getPackagePrefixesToCover() {
+        return "com.mongo.count";
+    }
+}