Added more test cases

Author: suddendust

document-store/src/test/java/org/hypertrace/core/documentstore/JSONDocumentTest.java

@@ -2,6 +2,7 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.IOException;
 import java.util.Map;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
@@ -106,4 +107,36 @@ public void testToStringMethod() throws Exception {
     // toString should return the same as toJson
     Assertions.assertEquals(document.toJson(), document.toString());
   }
+
+  @Test
+  public void testJsonNodeConstructors() throws Exception {
+    JsonNode node = mapper.createObjectNode().put("test", "value");
+
+    // Test JsonNode constructor without DocumentType
+    JSONDocument doc1 = new JSONDocument(node);
+    Assertions.assertEquals(DocumentType.NESTED, doc1.getDocumentType());
+
+    // Test JsonNode constructor with DocumentType
+    JSONDocument doc2 = new JSONDocument(node, DocumentType.FLAT);
+    Assertions.assertEquals(DocumentType.FLAT, doc2.getDocumentType());
+  }
+
+  @Test
+  public void testStringConstructorWithInvalidJson() {
+    // Test invalid JSON handling
+    Assertions.assertThrows(
+        IOException.class,
+        () -> new JSONDocument("invalid json {"));
+
+    Assertions.assertThrows(
+        IOException.class,
+        () -> new JSONDocument("invalid json {", DocumentType.FLAT));
+  }
+
+  @Test
+  public void testNullHandling() throws Exception {
+    // Test null JsonNode
+    JSONDocument nullNodeDoc = new JSONDocument((JsonNode) null);
+    Assertions.assertNotNull(nullNodeDoc.toJson()); // Should handle gracefully
+  }
 }

document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java

@@ -36,6 +36,7 @@
 import java.util.UUID;
 import org.hypertrace.core.documentstore.CloseableIterator;
 import org.hypertrace.core.documentstore.Document;
+import org.hypertrace.core.documentstore.DocumentType;
 import org.hypertrace.core.documentstore.Filter;
 import org.hypertrace.core.documentstore.JSONDocument;
 import org.hypertrace.core.documentstore.Key;
@@ -186,6 +187,7 @@ void testUpdateAtomicWithFilter() throws IOException, SQLException {
 
     assertTrue(oldDocument.isPresent());
     assertEquals(document, oldDocument.get());
+    assertEquals(DocumentType.NESTED, document.getDocumentType());
 
     verify(mockClient, times(1)).getPooledConnection();
     verify(mockConnection, times(1)).prepareStatement(selectQuery);