Serialisation attempt

Signed-off-by: Andy Kwok <[email protected]>
opensearch-project · Nov 27, 2024 · 384764f · 384764f
1 parent d2e6a8c
commit 384764f
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 17 deletions.
diff --git a/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentRequest.java b/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentRequest.java
@@ -31,6 +31,8 @@
 @AllArgsConstructor
 public class IpEnrichmentRequest extends ActionRequest {
 
+    private static final String VERSION = "2.18.0.0";
+
     private String ipString;
 
     private String datasourceName;
@@ -87,7 +89,12 @@ public static IpEnrichmentRequest fromActionRequest(ActionRequest actionRequest)
         try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
             actionRequest.writeTo(osso);
             try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
-                return new IpEnrichmentRequest(input);
+                String objectVersion = input.readString();
+                if (VERSION.equals(objectVersion)) {
+                    return new IpEnrichmentRequest(input);
+                } else {
+                    throw new IllegalArgumentException("Fail to serialise IpEnrichmentRequest due to version mismatch");
+                }
             }
         } catch (IOException e) {
             throw new UncheckedIOException("Failed to parse ActionRequest into IpEnrichmentRequest", e);

diff --git a/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentResponse.java b/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentResponse.java
@@ -33,6 +33,8 @@
 @EqualsAndHashCode(callSuper = false)
 public class IpEnrichmentResponse extends ActionResponse {
 
+    private static final String VERSION = "2.18.0.0";
+
     private Map<String, Object> geoLocationData;
 
     /**
@@ -72,8 +74,15 @@ public static IpEnrichmentResponse fromActionResponse(ActionResponse actionRespo
         // Or else convert it
         try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
             actionResponse.writeTo(osso);
+
+
             try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
-                return new IpEnrichmentResponse(input);
+                String objectVersion = input.readString();
+                if (VERSION.equals(objectVersion)) {
+                    return new IpEnrichmentResponse(input);
+                } else {
+                    throw new IllegalArgumentException("Fail to serialise IpEnrichmentResponse due to version mismatch");
+                }
             }
         } catch (IOException e) {
             throw new UncheckedIOException("Failed to parse ActionResponse into IpEnrichmentResponse", e);

diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/IpEnrichmentTransportActionTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/IpEnrichmentTransportActionTests.java
@@ -53,21 +53,7 @@ public void testDoExecute_All_Succeed() {
 
         verify(listener, times(1)).onResponse(any(IpEnrichmentResponse.class));
     }
-
-    /**
-     * When dataSource is absent, but default is valid.
-     */
-    @Test
-    public void testDoExecute_WithDefaultDataSource() {
-        when(mockDataSource.getName()).thenReturn("defaultDataSourceName");
-        when(datasourceDao.getAllDatasources()).thenReturn(List.of(mockDataSource));
-        when(ip2GeoCachedDao.getIndexName(eq("defaultDataSourceName"))).thenReturn("defaultIndexName");
-        when(ip2GeoCachedDao.getGeoData(eq("defaultIndexName"), any())).thenReturn(Collections.emptyMap());
-
-        IpEnrichmentRequest request = new IpEnrichmentRequest("192.168.1.1", null);
-        action.doExecute(task, request, listener);
-        verify(listener, times(1)).onResponse(any(IpEnrichmentResponse.class));
-    }
+
 
     /**
      * No alternative dataSource, exception being thrown to indicate this.