Refactor network connection setting on Android

src/main/java/io/appium/java_client/android/AndroidDriver.java

@@ -25,6 +25,7 @@
 import io.appium.java_client.FindsByAndroidUIAutomator;
 import io.appium.java_client.LocksDevice;
 import io.appium.java_client.PressesKeyCode;
+import io.appium.java_client.android.connection.HasNetworkConnection;
 import io.appium.java_client.remote.MobilePlatform;
 import io.appium.java_client.screenrecording.CanRecordScreen;
 import io.appium.java_client.service.local.AppiumDriverLocalService;

src/main/java/io/appium/java_client/android/AndroidMobileCommandHelper.java

@@ -197,13 +197,12 @@ public class AndroidMobileCommandHelper extends MobileCommand {
     /**
      * This method forms a {@link Map} of parameters for the setting of device network connection.
      *
-     * @param connection The bitmask of the desired connection
+     * @param bitMask The bitmask of the desired connection
      * @return a key-value pair. The key is the command name. The value is a {@link Map} command arguments.
      */
-    public static Map.Entry<String, Map<String, ?>> setConnectionCommand(Connection connection) {
+    public static Map.Entry<String, Map<String, ?>> setConnectionCommand(long bitMask) {
         String[] parameters = new String[] {"name", "parameters"};
-        Object[] values =
-            new Object[] {"network_connection", ImmutableMap.of("type", connection.bitMask)};
+        Object[] values = new Object[] {"network_connection", ImmutableMap.of("type", bitMask)};
         return new AbstractMap.SimpleEntry<>(
                 SET_NETWORK_CONNECTION, prepareArguments(parameters, values));
     }

src/main/java/io/appium/java_client/android/Connection.java

@@ -18,7 +18,9 @@
 
 /**
  * for use with setting Network Connections on a mobile device.
+ * @deprecated Use {@link io.appium.java_client.android.connection.ConnectionState} instead
  */
+@Deprecated
 public enum Connection {
     NONE(0),
     AIRPLANE(1),
@@ -31,4 +33,8 @@ public enum Connection {
     Connection(int bitMask) {
         this.bitMask = bitMask;
     }
+
+    public int getBitMask() {
+        return bitMask;
+    }
 }

src/main/java/io/appium/java_client/android/connection/ConnectionState.java

@@ -0,0 +1,55 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * See the NOTICE file distributed with this work for additional
+ * information regarding copyright ownership.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.appium.java_client.android.connection;
+
+public class ConnectionState {
+    public static final long AIRPLANE_MODE_MASK = 0b001;
+    public static final long WIFI_MASK = 0b010;
+    public static final long DATA_MASK = 0b100;
+
+    private final long bitMask;
+
+    public long getBitMask() {
+        return bitMask;
+    }
+
+    public ConnectionState(long bitMask) {
+        this.bitMask = bitMask;
+    }
+
+    /**
+     * @return true if airplane mode is enabled.
+     */
+    public boolean isAirplaneModeEnabled() {
+        return (bitMask & AIRPLANE_MODE_MASK) != 0;
+    }
+
+    /**
+     * @return true if Wi-Fi connection is enabled.
+     */
+    public boolean isWiFiEnabled() {
+        return (bitMask & WIFI_MASK) != 0;
+    }
+
+    /**
+     * @return true if data connection is enabled.
+     */
+    public boolean isDataEnabled() {
+        return (bitMask & DATA_MASK) != 0;
+    }
+}
+

src/main/java/io/appium/java_client/android/connection/ConnectionStateBuilder.java

@@ -0,0 +1,127 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * See the NOTICE file distributed with this work for additional
+ * information regarding copyright ownership.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.appium.java_client.android.connection;
+
+import static io.appium.java_client.android.connection.ConnectionState.AIRPLANE_MODE_MASK;
+import static io.appium.java_client.android.connection.ConnectionState.DATA_MASK;
+import static io.appium.java_client.android.connection.ConnectionState.WIFI_MASK;
+
+public class ConnectionStateBuilder {
+    private long bitMask;
+
+    /**
+     * Initializes connection state builder with the default value (all off).
+     */
+    public ConnectionStateBuilder() {
+        this.bitMask = 0;
+    }
+
+    /**
+     * Initializes connection state builder with the the predefined bit mask.
+     * This constructor might be handy to change an existing connection state.
+     *
+     * @param bitMask the actual initial state bit mask to set
+     */
+    public ConnectionStateBuilder(long bitMask) {
+        this.bitMask = bitMask;
+    }
+
+    /**
+     * Initializes connection state builder with the the predefined bit mask.
+     * This constructor might be handy to change an existing connection state.
+     *
+     * @param state the actual initial state to set
+     */
+    public ConnectionStateBuilder(ConnectionState state) {
+        this(state.getBitMask());
+    }
+
+    /**
+     * Sets airplane mode to enabled state if it was disabled.
+     * This option only works up to Android 6.
+     * Enabling the airplane mode on the device will automatically
+     * disable Wi-Fi and data connections.
+     *
+     * @return self instance for chaining
+     */
+    public ConnectionStateBuilder withAirplaneModeEnabled() {
+        bitMask |= AIRPLANE_MODE_MASK;
+        return this;
+    }
+
+    /**
+     * Sets airplane mode to disabled state if it was enabled.
+     * This option only works up to Android 6.
+     *
+     * @return self instance for chaining
+     */
+    public ConnectionStateBuilder withAirplaneModeDisabled() {
+        bitMask &= ~AIRPLANE_MODE_MASK;
+        return this;
+    }
+
+    /**
+     * Sets Wi-Fi connection mode to enabled state if it was disabled.
+     *
+     * @return self instance for chaining
+     */
+    public ConnectionStateBuilder withWiFiEnabled() {
+        bitMask |= WIFI_MASK;
+        return this;
+    }
+
+    /**
+     * Sets Wi-Fi connection mode to disabled state if it was enabled.
+     *
+     * @return self instance for chaining
+     */
+    public ConnectionStateBuilder withWiFiDisabled() {
+        bitMask &= ~WIFI_MASK;
+        return this;
+    }
+
+    /**
+     * Sets data connection mode to enabled state if it was disabled.
+     * This option only works on rooted devices or on emulators.
+     *
+     * @return self instance for chaining
+     */
+    public ConnectionStateBuilder withDataEnabled() {
+        bitMask |= DATA_MASK;
+        return this;
+    }
+
+    /**
+     * Sets data connection mode to disabled state if it was enabled.
+     * This option only works on rooted devices or on emulators.
+     *
+     * @return self instance for chaining
+     */
+    public ConnectionStateBuilder withDataDisabled() {
+        bitMask &= ~DATA_MASK;
+        return this;
+    }
+
+    /**
+     * Builds connection state instance, which is ready to be passed as Appium server parameter.
+     *
+     * @return ConnectionState instance
+     */
+    public ConnectionState build() {
+        return new ConnectionState(bitMask);
+    }
+}

src/main/java/io/appium/java_client/android/HasNetworkConnection.java → src/main/java/io/appium/java_client/android/connection/HasNetworkConnection.java

@@ -14,43 +14,44 @@
  * limitations under the License.
  */
 
-package io.appium.java_client.android;
+package io.appium.java_client.android.connection;
 
 import static io.appium.java_client.android.AndroidMobileCommandHelper.getNetworkConnectionCommand;
 import static io.appium.java_client.android.AndroidMobileCommandHelper.setConnectionCommand;
 
 import io.appium.java_client.CommandExecutionHelper;
 import io.appium.java_client.ExecutesMethod;
-import org.openqa.selenium.WebDriverException;
+import io.appium.java_client.android.Connection;
 
 public interface HasNetworkConnection extends ExecutesMethod {
-
     /**
      * Set the network connection of the device.
      *
      * @param connection The bitmask of the desired connection
+     * @deprecated use {@link #setConnection(ConnectionState)} instead
      */
+    @Deprecated
     default void setConnection(Connection connection) {
-        CommandExecutionHelper.execute(this, setConnectionCommand(connection));
+        CommandExecutionHelper.execute(this, setConnectionCommand(connection.getBitMask()));
     }
 
+    /**
+     * Set the network connection of the device.
+     *
+     * @param connection The bitmask of the desired connection
+     * @return Connection object, which represents the resulting state
+     */
+    default ConnectionState setConnection(ConnectionState connection) {
+        return new ConnectionState(CommandExecutionHelper.execute(this,
+                setConnectionCommand(connection.getBitMask())));
+    }
 
     /**
      * Get the current network settings of the device.
      *
-     * @return Connection object will let you inspect the status
-     *     of None, AirplaneMode, Wifi, Data and All connections
+     * @return Connection object, which lets you to inspect the current status
      */
-    default Connection getConnection() {
-        long bitMask = CommandExecutionHelper.execute(this, getNetworkConnectionCommand());
-        Connection[] types = Connection.values();
-
-        for (Connection connection: types) {
-            if (connection.bitMask == bitMask) {
-                return connection;
-            }
-        }
-        throw new WebDriverException("The unknown network connection "
-            + "type has been returned. The bitmask is " + bitMask);
+    default ConnectionState getConnection() {
+        return new ConnectionState(CommandExecutionHelper.execute(this, getNetworkConnectionCommand()));
     }
 }

src/test/java/io/appium/java_client/android/AndroidConnectionTest.java

@@ -16,33 +16,50 @@
 
 package io.appium.java_client.android;
 
-import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
+import io.appium.java_client.android.connection.ConnectionState;
+import io.appium.java_client.android.connection.ConnectionStateBuilder;
 import org.junit.FixMethodOrder;
 import org.junit.Test;
 import org.junit.runners.MethodSorters;
 
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 public class AndroidConnectionTest extends BaseAndroidTest {
 
-    @Test public void test1() {
-        driver.setConnection(Connection.WIFI);
-        assertEquals(Connection.WIFI,
-            driver.getConnection());
+    @Test
+    public void test1() {
+        ConnectionState state = driver.setConnection(new ConnectionStateBuilder()
+                .withWiFiEnabled()
+                .build());
+        assertTrue(state.isWiFiEnabled());
     }
 
-    @Test public void test2() {
-        driver.setConnection(Connection.NONE);
-        assertEquals(Connection.NONE,
-            driver.getConnection());
-        driver.setConnection(Connection.AIRPLANE);
-        assertEquals(Connection.AIRPLANE,
-            driver.getConnection());
+    @Test
+    public void test2() {
+        ConnectionState state = driver.setConnection(new ConnectionStateBuilder()
+                .withAirplaneModeDisabled()
+                .build());
+        assertFalse(state.isAirplaneModeEnabled());
+        assertFalse(state.isWiFiEnabled());
+        assertFalse(state.isDataEnabled());
+        state = driver.setConnection(new ConnectionStateBuilder(state)
+                .withAirplaneModeEnabled()
+                .build());
+        assertTrue(state.isAirplaneModeEnabled());
     }
 
-    @Test public void test3() {
-        driver.setConnection(Connection.ALL);
-        assertEquals(Connection.ALL,
-            driver.getConnection());
+    @Test
+    public void test3() {
+        ConnectionState state = driver.setConnection(
+                new ConnectionStateBuilder(driver.getConnection())
+                        .withAirplaneModeDisabled()
+                        .withWiFiEnabled()
+                        .withDataEnabled()
+                        .build());
+        assertFalse(state.isAirplaneModeEnabled());
+        assertTrue(state.isWiFiEnabled());
+        assertTrue(state.isDataEnabled());
     }
 }