diff --git a/rag/rag-springai-ollama-llm/src/main/java/com/learning/ai/llmragwithspringai/controller/DataIndexController.java b/rag/rag-springai-ollama-llm/src/main/java/com/learning/ai/llmragwithspringai/controller/DataIndexController.java
index 332c773..e24099e 100644
--- a/rag/rag-springai-ollama-llm/src/main/java/com/learning/ai/llmragwithspringai/controller/DataIndexController.java
+++ b/rag/rag-springai-ollama-llm/src/main/java/com/learning/ai/llmragwithspringai/controller/DataIndexController.java
@@ -1,6 +1,7 @@
 package com.learning.ai.llmragwithspringai.controller;
 
 import com.learning.ai.llmragwithspringai.service.DataIndexerService;
+import java.util.Map;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
@@ -29,7 +30,7 @@ public ResponseEntity<String> load(@RequestPart("file") MultipartFile multipartF
     }
 
     @GetMapping("count")
-    public long count() {
-        return dataIndexerService.count();
+    public Map<String, Long> count() {
+        return Map.of("count", dataIndexerService.count());
     }
 }
diff --git a/rag/rag-springai-ollama-llm/src/test/java/com/learning/ai/llmragwithspringai/LlmRagWithSpringAiApplicationIntTest.java b/rag/rag-springai-ollama-llm/src/test/java/com/learning/ai/llmragwithspringai/LlmRagWithSpringAiApplicationIntTest.java
index 5e2642c..dd246e1 100644
--- a/rag/rag-springai-ollama-llm/src/test/java/com/learning/ai/llmragwithspringai/LlmRagWithSpringAiApplicationIntTest.java
+++ b/rag/rag-springai-ollama-llm/src/test/java/com/learning/ai/llmragwithspringai/LlmRagWithSpringAiApplicationIntTest.java
@@ -39,11 +39,21 @@ void uploadPdfContent() throws IOException {
                 .when()
                 .post("/api/data/v1/upload")
                 .then()
-                .statusCode(200);
+                .statusCode(200)
+                .log()
+                .all();
     }
 
-    private File getFile(String fileName) throws IOException {
-        return new ClassPathResource(fileName).getFile();
+    @Test
+    @Order(2)
+    void uploadPdfContentCount() {
+        given().when()
+                .get("/api/data/v1/count")
+                .then()
+                .statusCode(200)
+                .body("count", is(1))
+                .log()
+                .all();
     }
 
     @Test
@@ -60,6 +70,20 @@ void testRag() {
                 .all();
     }
 
+    @Test
+    @Order(102)
+    void testRag2() {
+        given().contentType(ContentType.JSON)
+                .body(new AIChatRequest("Did Rohit Sharma won ICC Mens T20 World Cup 2007 ?"))
+                .when()
+                .post("/api/ai/chat")
+                .then()
+                .statusCode(200)
+                .body("queryResponse", containsString("Yes"))
+                .log()
+                .all();
+    }
+
     @Test
     @Order(111)
     void testEmptyQuery() {
@@ -115,4 +139,8 @@ void testSpecialCharactersInQuery() {
                 .body("violations[0].message", containsString("Invalid characters in query"))
                 .log();
     }
+
+    private File getFile(String fileName) throws IOException {
+        return new ClassPathResource(fileName).getFile();
+    }
 }