From ee8bac23923d0aa39bd4a73c52d0a50823a5249c Mon Sep 17 00:00:00 2001
From: Bartosz Witkowski <bartosz.witkowski@like-a-boss.net>
Date: Fri, 14 Apr 2023 16:59:36 +0200
Subject: [PATCH] Added x-correlation-id to proto requests

---
 .../sdk/dto/requests/ProtoRequest.java        |  3 ++
 .../TimeFilterFastProtoDistanceRequest.java   | 32 +++++++++++++++++++
 .../requests/TimeFilterFastProtoRequest.java  | 11 +++++++
 .../sdk/TimeFilterFastProtoTest.java          |  2 +-
 4 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/traveltime/sdk/dto/requests/ProtoRequest.java b/src/main/java/com/traveltime/sdk/dto/requests/ProtoRequest.java
index f284e12..cea0096 100644
--- a/src/main/java/com/traveltime/sdk/dto/requests/ProtoRequest.java
+++ b/src/main/java/com/traveltime/sdk/dto/requests/ProtoRequest.java
@@ -28,6 +28,8 @@ public abstract class ProtoRequest<T> {
 
     public abstract Either<TravelTimeError, T> parseBytes(byte[] body);
 
+    public abstract String getCorrelationId();
+
     protected Either<TravelTimeError, TimeFilterFastResponseOuterClass.TimeFilterFastResponse> getProtoResponse(byte[] body) {
         return Try
             .of(() -> TimeFilterFastResponseOuterClass.TimeFilterFastResponse.parseFrom(body))
@@ -45,6 +47,7 @@ protected Request createProtobufRequest(
             .headers(credentials.getBasicCredentialsHeaders())
             .addHeader("Content-Type", AcceptType.APPLICATION_OCTET_STREAM.getValue())
             .addHeader("User-Agent", "Travel Time Java SDK " + Version.getVersion())
+            .addHeader("X-Correlation-ID", getCorrelationId())
             .post(RequestBody.create(requestBody))
             .build();
     }
diff --git a/src/main/java/com/traveltime/sdk/dto/requests/TimeFilterFastProtoDistanceRequest.java b/src/main/java/com/traveltime/sdk/dto/requests/TimeFilterFastProtoDistanceRequest.java
index a058420..e99a5f8 100644
--- a/src/main/java/com/traveltime/sdk/dto/requests/TimeFilterFastProtoDistanceRequest.java
+++ b/src/main/java/com/traveltime/sdk/dto/requests/TimeFilterFastProtoDistanceRequest.java
@@ -39,6 +39,8 @@ public class TimeFilterFastProtoDistanceRequest extends ProtoRequest<TimeFilterF
     @NonNull
     Integer travelTime;
 
+    String correlationId;
+
     private byte[] createByteArray() {
 
         RequestsCommon.Coords departure = RequestsCommon
@@ -145,6 +147,15 @@ public TimeFilterFastProtoDistanceResponse merge(List<TimeFilterFastProtoDistanc
         return new TimeFilterFastProtoDistanceResponse(times, distances);
     }
 
+    @Override
+    public String getCorrelationId() {
+        if (correlationId == null) {
+            return "no-x-correlation-id";
+        } else {
+            return correlationId;
+        }
+    }
+
     /**
      * @param originCoordinate       The coordinates of location we should start the search from.
      * @param destinationCoordinates The coordinates of locations we run the search to. If the class implementing this list
@@ -167,5 +178,26 @@ public TimeFilterFastProtoDistanceRequest(
         this.travelTime = travelTime;
         this.transportation = transportation;
         this.country = country;
+        this.correlationId = null;
+    }
+
+    public TimeFilterFastProtoDistanceRequest(
+        @NonNull Coordinates originCoordinate,
+        @NonNull List<Coordinates> destinationCoordinates,
+        @NonNull Integer travelTime,
+        @NonNull Transportation transportation,
+        @NonNull Country country,
+        @NonNull String correlationId
+    ) {
+        this.originCoordinate = originCoordinate;
+        if (destinationCoordinates instanceof RandomAccess) {
+            this.destinationCoordinates = destinationCoordinates;
+        } else {
+            this.destinationCoordinates = new ArrayList<>(destinationCoordinates);
+        }
+        this.travelTime = travelTime;
+        this.transportation = transportation;
+        this.country = country;
+        this.correlationId = correlationId;
     }
 }
diff --git a/src/main/java/com/traveltime/sdk/dto/requests/TimeFilterFastProtoRequest.java b/src/main/java/com/traveltime/sdk/dto/requests/TimeFilterFastProtoRequest.java
index 8df26e1..7222bcb 100644
--- a/src/main/java/com/traveltime/sdk/dto/requests/TimeFilterFastProtoRequest.java
+++ b/src/main/java/com/traveltime/sdk/dto/requests/TimeFilterFastProtoRequest.java
@@ -28,6 +28,8 @@ public class TimeFilterFastProtoRequest extends ProtoRequest<TimeFilterFastProto
     @NonNull
     OneToMany oneToMany;
 
+    String correlationId;
+
     private byte[] createByteArray() {
         Coordinates origin = oneToMany.getOriginCoordinate();
 
@@ -125,6 +127,15 @@ private Either<TravelTimeError, TimeFilterFastProtoResponse> parseResponse(
             return Either.right(new TimeFilterFastProtoResponse(response.getProperties().getTravelTimesList()));
     }
 
+    @Override
+    public String getCorrelationId() {
+        if (correlationId == null) {
+            return "no-x-correlation-id";
+        } else {
+            return correlationId;
+        }
+    }
+
     @Override
     public Either<TravelTimeError, Request> createRequest(HttpUrl baseUri, TravelTimeCredentials credentials) {
         String countryCode = oneToMany.getCountry().getValue();
diff --git a/src/test/java/com/traveltime/sdk/TimeFilterFastProtoTest.java b/src/test/java/com/traveltime/sdk/TimeFilterFastProtoTest.java
index 98413e9..d01cf5d 100644
--- a/src/test/java/com/traveltime/sdk/TimeFilterFastProtoTest.java
+++ b/src/test/java/com/traveltime/sdk/TimeFilterFastProtoTest.java
@@ -161,6 +161,6 @@ public TimeFilterFastProtoRequest oneToMany(
             Country.UNITED_KINGDOM
         );
 
-        return new TimeFilterFastProtoRequest(oneToMany);
+        return new TimeFilterFastProtoRequest(oneToMany, null);
     }
 }