From 727e7bd171a32fef29e89b5152963b893cd95b39 Mon Sep 17 00:00:00 2001
From: Stefan Kalscheuer <stefan@stklcode.de>
Date: Tue, 27 Dec 2016 17:29:59 +0100
Subject: [PATCH] Response types checked

---
 .../de/stklcode/pubtrans/ura/UraClient.java   | 25 +++++++++----------
 .../stklcode/pubtrans/ura/UraClientTest.java  |  5 ++++
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/src/main/java/de/stklcode/pubtrans/ura/UraClient.java b/src/main/java/de/stklcode/pubtrans/ura/UraClient.java
index d903f86..80fcfd2 100644
--- a/src/main/java/de/stklcode/pubtrans/ura/UraClient.java
+++ b/src/main/java/de/stklcode/pubtrans/ura/UraClient.java
@@ -40,6 +40,9 @@ public class UraClient {
     private static final String PAR_TRIP_ID = "TripID";
     private static final String PAR_ESTTIME = "EstimatedTime";
 
+    private static final Integer RES_TYPE_STOP = 0;
+    private static final Integer RES_TYPE_PREDICTION = 1;
+
 
     private static final String[] REQUEST_STOP = {PAR_STOP_NAME, PAR_STOP_ID, PAR_STOP_INDICATOR, PAR_STOP_STATE, PAR_GEOLOCATION};
     private static final String[] REQUEST_TRIP = {PAR_STOP_NAME, PAR_STOP_ID, PAR_STOP_INDICATOR, PAR_STOP_STATE, PAR_GEOLOCATION,
@@ -142,13 +145,11 @@ public class UraClient {
         try (InputStream is = requestInstant(REQUEST_TRIP, stops, lines);
              BufferedReader br = new BufferedReader(new InputStreamReader(is))) {
             String line;
-            boolean first = false;
-            while ((line = br.readLine()) != null) {
-                if (!first) {
-                    first = true;
-                    continue;
-                }
-                trips.add(new Trip(mapper.readValue(line, List.class)));
+            while ((line = br.readLine()) != null && (limit == null || trips.size() < limit)) {
+                List l = mapper.readValue(line, List.class);
+                /* Check if result exists and has correct response type */
+                if (l != null && l.size() > 0 && l.get(0).equals(RES_TYPE_PREDICTION))
+                    trips.add(new Trip(l));
             }
         } catch (IOException e) {
             e.printStackTrace();
@@ -166,13 +167,11 @@ public class UraClient {
         try (InputStream is = requestInstant(REQUEST_STOP, null, null);
              BufferedReader br = new BufferedReader(new InputStreamReader(is))) {
             String line;
-            boolean first = false;
             while ((line = br.readLine()) != null) {
-                if (!first) {
-                    first = true;
-                    continue;
-                }
-                stops.add(new Stop(mapper.readValue(line, List.class)));
+                List l = mapper.readValue(line, List.class);
+                /* Check if result exists and has correct response type */
+                if (l != null && l.size() > 0 && l.get(0).equals(RES_TYPE_STOP))
+                    stops.add(new Stop(l));
             }
         } catch (IOException e) {
             e.printStackTrace();
diff --git a/src/test/java/de/stklcode/pubtrans/ura/UraClientTest.java b/src/test/java/de/stklcode/pubtrans/ura/UraClientTest.java
index 56a446a..e3c2917 100644
--- a/src/test/java/de/stklcode/pubtrans/ura/UraClientTest.java
+++ b/src/test/java/de/stklcode/pubtrans/ura/UraClientTest.java
@@ -63,6 +63,11 @@ public class UraClientTest {
         assertThat(trips.get(7).getEstimatedTime(), is(1482854580000L));
         assertThat(trips.get(8).getVisitID(), is(30));
         assertThat(trips.get(9).getStop().getId(), is("100002"));
+
+        /* Get limited number of trips */
+        PowerMockito.when(mockURL.openStream()).thenReturn(getClass().getResourceAsStream("instant_trips_all.txt"));
+        trips = new UraClient("mocked").getTrips(5);
+        assertThat(trips, hasSize(5));
     }
 
     @Test