calculate and show gateway load related to transport

9 files changed, 165 insertions, 93 deletions

diff --git a/app/src/main/java/se/leap/bitmaskclient/base/fragments/EipFragment.java b/app/src/main/java/se/leap/bitmaskclient/base/fragments/EipFragment.java
index f8053f5e..e9300584 100644
--- a/app/src/main/java/se/leap/bitmaskclient/base/fragments/EipFragment.java
+++ b/app/src/main/java/se/leap/bitmaskclient/base/fragments/EipFragment.java
@@ -56,12 +56,14 @@ import de.blinkt.openvpn.core.ConnectionStatus;
 import de.blinkt.openvpn.core.IOpenVPNServiceInternal;
 import de.blinkt.openvpn.core.OpenVPNService;
 import de.blinkt.openvpn.core.VpnStatus;
+import de.blinkt.openvpn.core.connection.Connection;
 import se.leap.bitmaskclient.BuildConfig;
 import se.leap.bitmaskclient.R;
 import se.leap.bitmaskclient.base.FragmentManagerEnhanced;
 import se.leap.bitmaskclient.base.MainActivity;
 import se.leap.bitmaskclient.base.models.Provider;
 import se.leap.bitmaskclient.base.models.ProviderObservable;
+import se.leap.bitmaskclient.base.utils.PreferenceHelper;
 import se.leap.bitmaskclient.base.views.LocationButton;
 import se.leap.bitmaskclient.base.views.MainButton;
 import se.leap.bitmaskclient.eip.EipCommand;
@@ -426,7 +428,8 @@ public class EipFragment extends Fragment implements Observer {
         } else if (eipStatus.isConnected()) {
             setMainButtonEnabled(true);
             mainButton.updateState(true, false, false);
-            locationButton.setLocationLoad(gatewaysManager.getLoadForLocation(VpnStatus.getLastConnectedVpnName()));
+            Connection.TransportType transportType = PreferenceHelper.getUseBridges(getContext()) ? Connection.TransportType.OBFS4 : Connection.TransportType.OPENVPN;
+            locationButton.setLocationLoad(gatewaysManager.getLoadForLocation(VpnStatus.getLastConnectedVpnName(), transportType));
             locationButton.setText(VpnStatus.getLastConnectedVpnName());
             locationButton.showBridgeIndicator(VpnStatus.isUsingBridges());
             locationButton.showRecommendedIndicator(getPreferredCity(getContext())== null);
diff --git a/app/src/main/java/se/leap/bitmaskclient/base/fragments/GatewaySelectionFragment.java b/app/src/main/java/se/leap/bitmaskclient/base/fragments/GatewaySelectionFragment.java
index 12aedb12..51ebe359 100644
--- a/app/src/main/java/se/leap/bitmaskclient/base/fragments/GatewaySelectionFragment.java
+++ b/app/src/main/java/se/leap/bitmaskclient/base/fragments/GatewaySelectionFragment.java
@@ -19,6 +19,7 @@ package se.leap.bitmaskclient.base.fragments;
 import android.app.Activity;
 import android.content.Context;
 import android.content.Intent;
+import android.content.SharedPreferences;
 import android.os.Bundle;
 import android.util.Log;
 import android.view.LayoutInflater;
@@ -48,18 +49,21 @@ import se.leap.bitmaskclient.eip.EipCommand;
 import se.leap.bitmaskclient.eip.EipStatus;
 import se.leap.bitmaskclient.eip.GatewaysManager;
 
+import static android.content.Context.MODE_PRIVATE;
 import static android.view.View.VISIBLE;
 import static de.blinkt.openvpn.core.connection.Connection.TransportType.OBFS4;
 import static de.blinkt.openvpn.core.connection.Connection.TransportType.OPENVPN;
 import static se.leap.bitmaskclient.base.MainActivity.ACTION_SHOW_DIALOG_FRAGMENT;
 import static se.leap.bitmaskclient.base.MainActivity.ACTION_SHOW_VPN_FRAGMENT;
 import static se.leap.bitmaskclient.base.models.Constants.LOCATION;
+import static se.leap.bitmaskclient.base.models.Constants.SHARED_PREFERENCES;
+import static se.leap.bitmaskclient.base.models.Constants.USE_BRIDGES;
 
 interface LocationListSelectionListener {
     void onLocationManuallySelected(Location location);
 }
 
-public class GatewaySelectionFragment extends Fragment implements Observer, LocationListSelectionListener {
+public class GatewaySelectionFragment extends Fragment implements Observer, LocationListSelectionListener, SharedPreferences.OnSharedPreferenceChangeListener {
 
     private static final String TAG = GatewaySelectionFragment.class.getSimpleName();
 
@@ -69,6 +73,8 @@ public class GatewaySelectionFragment extends Fragment implements Observer, Loca
     private SelectLocationEntry recommendedLocation;
     private GatewaysManager gatewaysManager;
     private EipStatus eipStatus;
+    private SharedPreferences preferences;
+    private Connection.TransportType selectedTransport;
 
     public GatewaySelectionFragment() {
         // Required empty public constructor
@@ -80,6 +86,9 @@ public class GatewaySelectionFragment extends Fragment implements Observer, Loca
         gatewaysManager = new GatewaysManager(getContext());
         eipStatus = EipStatus.getInstance();
         eipStatus.addObserver(this);
+        preferences = getContext().getSharedPreferences(SHARED_PREFERENCES, MODE_PRIVATE);
+        selectedTransport = PreferenceHelper.getUseBridges(preferences) ? OBFS4 : OPENVPN;
+        preferences.registerOnSharedPreferenceChangeListener(this);
     }
 
     @Override
@@ -100,6 +109,7 @@ public class GatewaySelectionFragment extends Fragment implements Observer, Loca
     public void onDestroyView() {
         super.onDestroyView();
         eipStatus.deleteObserver(this);
+        preferences.unregisterOnSharedPreferenceChangeListener(this);
     }
 
     private void initRecyclerView() {
@@ -107,7 +117,7 @@ public class GatewaySelectionFragment extends Fragment implements Observer, Loca
         recyclerView.setHasFixedSize(true);
         LinearLayoutManager layoutManager = new LinearLayoutManager(this.getContext());
         recyclerView.setLayoutManager(layoutManager);
-        locationListAdapter = new LocationListAdapter(gatewaysManager.getGatewayLocations(), this);
+        locationListAdapter = new LocationListAdapter(gatewaysManager.getGatewayLocations(), this, selectedTransport);
         recyclerView.setAdapter(locationListAdapter);
         recyclerView.setVisibility(VISIBLE);
     }
@@ -135,7 +145,10 @@ public class GatewaySelectionFragment extends Fragment implements Observer, Loca
             }
         }
         location.selected = !isManualSelection;
-        recommendedLocation.setLocation(location);
+        if (!isManualSelection) {
+            locationListAdapter.unselectAll();
+        }
+        recommendedLocation.setLocation(location, selectedTransport);
     }
 
     protected void startEipService(String preferredCity) {
@@ -156,11 +169,11 @@ public class GatewaySelectionFragment extends Fragment implements Observer, Loca
     @Override
     public void onLocationManuallySelected(Location location) {
         recommendedLocation.setSelected(false);
-        String name = location.name;
-        Connection.TransportType selectedTransport = PreferenceHelper.getUseBridges(getContext()) ? OBFS4 : OPENVPN;
-        if (location.supportedTransports.contains(selectedTransport)) {
+        String name = location.getName();
+        if (location.supportsTransport(selectedTransport)) {
             startEipService(name);
         } else {
+            locationListAdapter.unselectAll();
             askToChangeTransport(name);
         }
     }
@@ -186,8 +199,17 @@ public class GatewaySelectionFragment extends Fragment implements Observer, Loca
         }
     }
 
+    @Override
+    public void onSharedPreferenceChanged(SharedPreferences sharedPreferences, String key) {
+        if (key.equals(USE_BRIDGES)) {
+            selectedTransport = PreferenceHelper.getUseBridges(sharedPreferences) ? OBFS4 : OPENVPN;
+            locationListAdapter.updateTransport(selectedTransport);
+        }
+    }
+
     static class LocationListAdapter extends RecyclerView.Adapter<LocationListAdapter.ViewHolder> {
         private static final String TAG = LocationListAdapter.class.getSimpleName();
+        private Connection.TransportType transport;
         private final List<Location> values;
         private final WeakReference<LocationListSelectionListener> callback;
 
@@ -210,6 +232,11 @@ public class GatewaySelectionFragment extends Fragment implements Observer, Loca
             notifyItemRemoved(position);
         }
 
+        public void updateTransport(Connection.TransportType transportType) {
+            transport = transportType;
+            notifyDataSetChanged();
+        }
+
         public void unselectAll() {
             for (Location l : values) {
                 l.selected = false;
@@ -217,9 +244,10 @@ public class GatewaySelectionFragment extends Fragment implements Observer, Loca
             notifyDataSetChanged();
         }
 
-        public LocationListAdapter(List<Location> data, LocationListSelectionListener selectionListener) {
+        public LocationListAdapter(List<Location> data, LocationListSelectionListener selectionListener, Connection.TransportType selectedTransport) {
             values = data;
             callback = new WeakReference<>(selectionListener);
+            transport = selectedTransport;
         }
 
         @NonNull
@@ -234,7 +262,7 @@ public class GatewaySelectionFragment extends Fragment implements Observer, Loca
         @Override
         public void onBindViewHolder(ViewHolder holder, final int position) {
             final Location location = values.get(position);
-            holder.entry.setLocation(location);
+            holder.entry.setLocation(location, transport);
             holder.entry.setOnClickListener(v -> {
                 Log.d(TAG, "onClick view at position clicked: " + position);
                 LocationListSelectionListener listener = callback.get();
diff --git a/app/src/main/java/se/leap/bitmaskclient/base/models/Location.java b/app/src/main/java/se/leap/bitmaskclient/base/models/Location.java
index 599a358a..e0ce9e8b 100644
--- a/app/src/main/java/se/leap/bitmaskclient/base/models/Location.java
+++ b/app/src/main/java/se/leap/bitmaskclient/base/models/Location.java
@@ -19,29 +19,62 @@ package se.leap.bitmaskclient.base.models;
 
 import androidx.annotation.NonNull;
 
+import java.util.HashMap;
 import java.util.HashSet;
 
 import de.blinkt.openvpn.core.connection.Connection;
 
 public class Location implements Cloneable {
-    @NonNull public String name = "";
-    @NonNull public HashSet<Connection.TransportType> supportedTransports = new HashSet<>();
-    public double averageLoad;
-    public int numberOfGateways;
+    @NonNull private String name = "";
+    @NonNull private HashMap<Connection.TransportType, Double> averageLoad = new HashMap<>();
+    @NonNull private HashMap<Connection.TransportType, Integer> numberOfGateways = new HashMap<>();
     public boolean selected;
 
     public Location() {}
 
-    public Location(@NonNull String name, double averageLoad, int numberOfGateways, @NonNull HashSet<Connection.TransportType> supportedTransports, boolean selected) {
+    public Location(@NonNull String name,
+                    @NonNull HashMap<Connection.TransportType, Double> averageLoad,
+                    @NonNull HashMap<Connection.TransportType, Integer> numberOfGateways,
+                    boolean selected) {
         this.name = name;
         this.averageLoad = averageLoad;
         this.numberOfGateways = numberOfGateways;
-        this.supportedTransports = supportedTransports;
         this.selected = selected;
     }
 
     public boolean hasLocationInfo() {
-        return numberOfGateways != 0 && supportedTransports.size() != 0 && !name.isEmpty();
+        return !numberOfGateways.isEmpty() && !averageLoad.isEmpty() && !name.isEmpty();
+    }
+
+    public boolean supportsTransport(Connection.TransportType transportType) {
+        return numberOfGateways.containsKey(transportType);
+    }
+
+    public void setAverageLoad(Connection.TransportType transportType, double load) {
+        averageLoad.put(transportType, load);
+    }
+
+    public double getAverageLoad(Connection.TransportType transportType) {
+        if (averageLoad.containsKey(transportType)) {
+            return averageLoad.get(transportType);
+        }
+        return 0;
+    }
+
+    public void setNumberOfGateways(Connection.TransportType transportType, int numbers) {
+        numberOfGateways.put(transportType, numbers);
+    }
+
+    public int getNumberOfGateways(Connection.TransportType transportType) {
+        if (numberOfGateways.containsKey(transportType)) {
+            return numberOfGateways.get(transportType);
+        }
+        return 0;
+    }
+
+    @NonNull
+    public String getName() {
+        return name;
     }
 
     @Override
@@ -51,21 +84,16 @@ public class Location implements Cloneable {
 
         Location location = (Location) o;
 
-        if (Double.compare(location.averageLoad, averageLoad) != 0) return false;
-        if (numberOfGateways != location.numberOfGateways) return false;
         if (!name.equals(location.name)) return false;
-        return supportedTransports.equals(location.supportedTransports);
+        if (!averageLoad.equals(location.averageLoad)) return false;
+        return numberOfGateways.equals(location.numberOfGateways);
     }
 
     @Override
     public int hashCode() {
-        int result;
-        long temp;
-        result = name.hashCode();
-        result = 31 * result + supportedTransports.hashCode();
-        temp = Double.doubleToLongBits(averageLoad);
-        result = 31 * result + (int) (temp ^ (temp >>> 32));
-        result = 31 * result + numberOfGateways;
+        int result = name.hashCode();
+        result = 31 * result + averageLoad.hashCode();
+        result = 31 * result + numberOfGateways.hashCode();
         return result;
     }
 
@@ -73,10 +101,8 @@ public class Location implements Cloneable {
     public Location clone() throws CloneNotSupportedException {
         Location copy = (Location) super.clone();
         copy.name = this.name;
-        copy.supportedTransports = (HashSet<Connection.TransportType>) this.supportedTransports.clone();
-        copy.numberOfGateways = this.numberOfGateways;
-        copy.averageLoad = this.averageLoad;
+        copy.numberOfGateways = (HashMap<Connection.TransportType, Integer>) this.numberOfGateways.clone();
+        copy.averageLoad = (HashMap<Connection.TransportType, Double>) this.averageLoad.clone();
         return copy;
     }
-
 }
diff --git a/app/src/main/java/se/leap/bitmaskclient/base/utils/ConfigHelper.java b/app/src/main/java/se/leap/bitmaskclient/base/utils/ConfigHelper.java
index 64b51960..92010992 100644
--- a/app/src/main/java/se/leap/bitmaskclient/base/utils/ConfigHelper.java
+++ b/app/src/main/java/se/leap/bitmaskclient/base/utils/ConfigHelper.java
@@ -197,6 +197,15 @@ public class ConfigHelper {
         return Calendar.getInstance().get(Calendar.ZONE_OFFSET) / 3600000;
     }
 
+    public static int timezoneDistance(int local_timezone, int remoteTimezone) {
+        // Distance along the numberline of Prime Meridian centric, assumes UTC-11 through UTC+12
+        int dist = Math.abs(local_timezone - remoteTimezone);
+        // Farther than 12 timezones and it's shorter around the "back"
+        if (dist > 12)
+            dist = 12 - (dist - 12); // Well i'll be. Absolute values make equations do funny things.
+        return dist;
+    }
+
     public static String getProviderFormattedString(Resources resources, @StringRes int resourceId) {
         String appName = resources.getString(R.string.app_name);
         return resources.getString(resourceId, appName);
diff --git a/app/src/main/java/se/leap/bitmaskclient/base/views/SelectLocationEntry.java b/app/src/main/java/se/leap/bitmaskclient/base/views/SelectLocationEntry.java
index f85df4ca..bf293a51 100644
--- a/app/src/main/java/se/leap/bitmaskclient/base/views/SelectLocationEntry.java
+++ b/app/src/main/java/se/leap/bitmaskclient/base/views/SelectLocationEntry.java
@@ -64,14 +64,14 @@ public class SelectLocationEntry extends RelativeLayout {
         title.setText(text);
         title.setVisibility(text != null ? VISIBLE : GONE);
     }
-    public void setLocation(Location location) {
+    public void setLocation(Location location, Connection.TransportType transportType) {
         boolean valid = location.hasLocationInfo();
         locationText.setVisibility(valid ? VISIBLE : GONE);
         locationIndicator.setVisibility(valid ? VISIBLE : GONE);
         bridgesView.setVisibility(valid ? VISIBLE : GONE);
-        locationText.setText(location.name);
-        locationIndicator.setLoad(Load.getLoadByValue(location.averageLoad));
-        bridgesView.setVisibility(location.supportedTransports.contains(Connection.TransportType.OBFS4) ? VISIBLE : GONE);
+        locationText.setText(location.getName());
+        locationIndicator.setLoad(Load.getLoadByValue(location.getAverageLoad(transportType)));
+        bridgesView.setVisibility(location.supportsTransport(Connection.TransportType.OBFS4) ? VISIBLE : GONE);
         selectedView.setChecked(location.selected);
     }
 
diff --git a/app/src/main/java/se/leap/bitmaskclient/eip/GatewaySelector.java b/app/src/main/java/se/leap/bitmaskclient/eip/GatewaySelector.java
index a48cc6d5..52030ce3 100644
--- a/app/src/main/java/se/leap/bitmaskclient/eip/GatewaySelector.java
+++ b/app/src/main/java/se/leap/bitmaskclient/eip/GatewaySelector.java
@@ -11,6 +11,7 @@ import java.util.Set;
 import java.util.TreeMap;
 
 import static se.leap.bitmaskclient.base.utils.ConfigHelper.getCurrentTimezone;
+import static se.leap.bitmaskclient.base.utils.ConfigHelper.timezoneDistance;
 
 public class GatewaySelector {
     private final static String TAG = GatewaySelector.class.getSimpleName();
@@ -67,12 +68,4 @@ public class GatewaySelector {
         return offsets;
     }
 
-    private int timezoneDistance(int local_timezone, int remote_timezone) {
-        // Distance along the numberline of Prime Meridian centric, assumes UTC-11 through UTC+12
-        int dist = Math.abs(local_timezone - remote_timezone);
-        // Farther than 12 timezones and it's shorter around the "back"
-        if (dist > 12)
-            dist = 12 - (dist - 12); // Well i'll be. Absolute values make equations do funny things.
-        return dist;
-    }
 }
diff --git a/app/src/main/java/se/leap/bitmaskclient/eip/GatewaysManager.java b/app/src/main/java/se/leap/bitmaskclient/eip/GatewaysManager.java
index 92a6af5a..a0867605 100644
--- a/app/src/main/java/se/leap/bitmaskclient/eip/GatewaysManager.java
+++ b/app/src/main/java/se/leap/bitmaskclient/eip/GatewaysManager.java
@@ -149,54 +149,63 @@ public class GatewaysManager {
 
             if (!locationNames.containsKey(name)) {
                 locationNames.put(name, locations.size());
-                // fake values for now
-                Random rand = new Random();
-                double averageLoad = rand.nextDouble(); //location.averageLoad;
-                Log.d(TAG, "getGatewayLocations - new averageLoad (" + name + " - " + gateway.getHost()+ "): " + averageLoad);
-
-                Location location = new Location(
-                        name,
-                        averageLoad
-                        /*gateway.getFullness()*/,
-                        1,
-                        gateway.getSupportedTransports(),
-                        name.equals(preferredCity)
-                        );
+                Location location = initLocation(name, gateway, preferredCity);
                 locations.add(location);
             } else {
                 int index = locationNames.get(gateway.getName());
                 Location location = locations.get(index);
-                location.averageLoad = (location.numberOfGateways * location.averageLoad + gateway.getFullness()) / (location.numberOfGateways + 1);
-                Log.d(TAG, "getGatewayLocations - updated averageLoad: (" + gateway.getName() + " - " + gateway.getHost()+ "): " + location.averageLoad);
-                location.numberOfGateways += 1;
-                location.supportedTransports.addAll(gateway.getSupportedTransports());
+                updateLocation(location, gateway, OBFS4);
+                updateLocation(location, gateway, OPENVPN);
                 locations.set(index, location);
             }
         }
-
         this.locations = locations;
         return locations;
     }
 
+    private Location initLocation(String name, Gateway gateway, String preferredCity) {
+        HashMap<Connection.TransportType, Double> averageLoadMap = new HashMap<>();
+        HashMap<Connection.TransportType, Integer> numberOfGatewaysMap = new HashMap<>();
+        if (gateway.getSupportedTransports().contains(OBFS4)) {
+            averageLoadMap.put(OBFS4, gateway.getFullness());
+            numberOfGatewaysMap.put(OBFS4, 1);
+        }
+        if (gateway.getSupportedTransports().contains(OPENVPN)) {
+            averageLoadMap.put(OPENVPN, gateway.getFullness());
+            numberOfGatewaysMap.put(OPENVPN, 1);
+        }
+        return new Location(
+                name,
+                averageLoadMap,
+                numberOfGatewaysMap,
+                name.equals(preferredCity));
+    }
+
+    private void updateLocation(Location location, Gateway gateway, Connection.TransportType transportType) {
+        if (gateway.getSupportedTransports().contains(transportType)) {
+            double averageLoad = location.getAverageLoad(transportType);
+            int numberOfGateways = location.getNumberOfGateways(transportType);
+            averageLoad = (numberOfGateways * averageLoad + gateway.getFullness()) / (numberOfGateways + 1);
+            numberOfGateways++;
+            location.setAverageLoad(transportType, averageLoad);
+            location.setNumberOfGateways(transportType, numberOfGateways);
+        }
+    }
+
     @Nullable
     public Location getLocation(String name) {
         List <Location> locations = getGatewayLocations();
         for (Location location : locations) {
-            if (location.name.equals(name)) {
+            if (location.getName().equals(name)) {
                 return location;
             }
         }
         return null;
     }
 
-    public Load getLoadForLocation(@Nullable String name) {
+    public Load getLoadForLocation(@Nullable String name, Connection.TransportType transportType) {
         Location location = getLocation(name);
-        if (location != null) {
-            return Load.getLoadByValue(location.averageLoad);
-        }
-
-        // location not found
-        return Load.UNKNOWN;
+        return Load.getLoadByValue(location.getAverageLoad(transportType));
     }
 
     private Gateway getGatewayFromTimezoneCalculation(int nClosest, Connection.TransportType transportType, @Nullable String city) {
diff --git a/app/src/test/java/se/leap/bitmaskclient/eip/GatewaySelectorTest.java b/app/src/test/java/se/leap/bitmaskclient/eip/GatewaySelectorTest.java
index 5186c462..6f26e494 100644
--- a/app/src/test/java/se/leap/bitmaskclient/eip/GatewaySelectorTest.java
+++ b/app/src/test/java/se/leap/bitmaskclient/eip/GatewaySelectorTest.java
@@ -27,6 +27,7 @@ import se.leap.bitmaskclient.base.utils.ConfigHelper;
 
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.powermock.api.mockito.PowerMockito.mockStatic;
 import static org.powermock.api.mockito.PowerMockito.when;
 import static se.leap.bitmaskclient.base.models.Constants.PROVIDER_PRIVATE_KEY;
@@ -75,6 +76,7 @@ public class GatewaySelectorTest {
     @Before
     public void setup() throws IOException, JSONException, ConfigParser.ConfigParseError {
         mockStatic(ConfigHelper.class);
+        when(ConfigHelper.timezoneDistance(anyInt(), anyInt())).thenCallRealMethod();
         mockTextUtils();
         eipDefinition = new JSONObject(getInputAsString(getClass().getClassLoader().getResourceAsStream("eip-service-four-gateways.json")));
         JSONArray gateways = eipDefinition.getJSONArray("gateways");
diff --git a/app/src/test/java/se/leap/bitmaskclient/eip/GatewaysManagerTest.java b/app/src/test/java/se/leap/bitmaskclient/eip/GatewaysManagerTest.java
index 27bfdb34..28311e18 100644
--- a/app/src/test/java/se/leap/bitmaskclient/eip/GatewaysManagerTest.java
+++ b/app/src/test/java/se/leap/bitmaskclient/eip/GatewaysManagerTest.java
@@ -23,27 +23,28 @@ import de.blinkt.openvpn.core.ConfigParser;
 import se.leap.bitmaskclient.base.models.Location;
 import se.leap.bitmaskclient.base.models.Provider;
 import se.leap.bitmaskclient.base.models.ProviderObservable;
+import se.leap.bitmaskclient.base.utils.ConfigHelper;
+import se.leap.bitmaskclient.base.utils.PreferenceHelper;
 import se.leap.bitmaskclient.testutils.MockHelper;
 import se.leap.bitmaskclient.testutils.MockSharedPreferences;
 import se.leap.bitmaskclient.testutils.TestSetupHelper;
-import se.leap.bitmaskclient.base.utils.ConfigHelper;
-import se.leap.bitmaskclient.base.utils.PreferenceHelper;
 
 import static de.blinkt.openvpn.core.connection.Connection.TransportType.OBFS4;
 import static de.blinkt.openvpn.core.connection.Connection.TransportType.OPENVPN;
 import static junit.framework.Assert.assertEquals;
-import static junit.framework.Assert.assertNotNull;
 import static junit.framework.Assert.assertNull;
-import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.powermock.api.mockito.PowerMockito.mockStatic;
 import static se.leap.bitmaskclient.base.models.Constants.GATEWAYS;
 import static se.leap.bitmaskclient.base.models.Constants.PROVIDER_EIP_DEFINITION;
 import static se.leap.bitmaskclient.base.models.Constants.PROVIDER_PRIVATE_KEY;
 import static se.leap.bitmaskclient.base.models.Constants.PROVIDER_VPN_CERTIFICATE;
+import static se.leap.bitmaskclient.base.models.Constants.USE_BRIDGES;
 import static se.leap.bitmaskclient.base.models.Provider.CA_CERT;
 import static se.leap.bitmaskclient.testutils.MockHelper.mockTextUtils;
 import static se.leap.bitmaskclient.testutils.TestSetupHelper.getProvider;
@@ -69,6 +70,8 @@ public class GatewaysManagerTest {
         mockTextUtils();
         when(ConfigHelper.getCurrentTimezone()).thenReturn(-1);
         when(ConfigHelper.stringEqual(anyString(), anyString())).thenCallRealMethod();
+        when(ConfigHelper.isIPv4(anyString())).thenCallRealMethod();
+        when(ConfigHelper.timezoneDistance(anyInt(), anyInt())).thenCallRealMethod();
         secrets = new JSONObject(getJsonStringFor("secrets.json"));
         sharedPreferences = new MockSharedPreferences();
         sharedPreferences.edit().
@@ -345,10 +348,7 @@ public class GatewaysManagerTest {
         assertNull(gatewaysManager.select(0, "Stockholm"));
     }
 
-    /*
-        This test is disabled since we fake load values in gateway manager with randomized numbers.
-     */
-   /* @Test
+    @Test
     public void testGetLocations_openvpn() {
         Provider provider = getProvider(null, null, null, null, null, null, "v4/riseup_eipservice_for_geoip_v4.json", "v4/riseup_geoip_v4.json");
 
@@ -360,43 +360,45 @@ public class GatewaysManagerTest {
 
         assertEquals(3, locations.size());
         for (Location location : locations) {
-            if ("Paris".equals(location.name)) {
-                assertEquals(3, location.numberOfGateways);
+            if ("Paris".equals(location.getName())) {
+                assertEquals(3, location.getNumberOfGateways(OPENVPN));
                 // manually calculate average load of paris gateways in "v4/riseup_geoip_v4.json"
                 double averageLoad = (0.3 + 0.36 + 0.92) / 3.0;
-                assertEquals(averageLoad, location.averageLoad);
+                assertEquals(averageLoad, location.getAverageLoad(OPENVPN));
             }
         }
-    } */
+    }
 
-    /*
-    This test is disabled since we fake load values in gateway manager with randomized numbers.
-    */
-    /*@Test
+    @Test
     public void testGetLocations_obfs4() {
         Provider provider = getProvider(null, null, null, null, null, null, "v4/riseup_eipservice_for_geoip_v4.json", "v4/riseup_geoip_v4.json");
 
         MockHelper.mockProviderObservable(provider);
         mockStatic(PreferenceHelper.class);
         when(PreferenceHelper.getUseBridges(any(Context.class))).thenReturn(true);
+        sharedPreferences.edit().putBoolean(USE_BRIDGES, true).commit();
         GatewaysManager gatewaysManager = new GatewaysManager(mockContext);
         List<Location> locations = gatewaysManager.getGatewayLocations();
 
-        assertEquals(2, locations.size());
+        assertEquals(3, locations.size());
         for (Location location : locations) {
-            if ("Montreal".equals(location.name)) {
-                assertEquals(1, location.numberOfGateways);
-                assertEquals(0.59, location.averageLoad);
+            if ("Montreal".equals(location.getName())) {
+                assertEquals(1, location.getNumberOfGateways(OBFS4));
+                assertEquals(0.59, location.getAverageLoad(OBFS4));
+                assertTrue(location.supportsTransport(OBFS4));
             }
-            if ("Paris".equals(location.name)) {
+            if ("Paris".equals(location.getName())) {
                 // checks that only gateways supporting obfs4 are taken into account
-                assertEquals(1, location.numberOfGateways);
-                assertEquals(0.36, location.averageLoad);
+                assertEquals(1, location.getNumberOfGateways(OBFS4));
+                assertEquals(0.36, location.getAverageLoad(OBFS4));
+                assertTrue(location.supportsTransport(OBFS4));
+            }
+            if ("Amsterdam".equals(location.getName())) {
+                assertFalse(location.supportsTransport(OBFS4));
             }
         }
 
     }
-     */
 
 
     private String getJsonStringFor(String filename) throws IOException {