After subscription enabled test wait for cellular network.

This is needed to reset the state of the device before the
end of the test.

Test: atest SubscriptionManagerTest
Bug: 158892113
Change-Id: Ib737d906bfe304e825f4498c63c9ee5b24cb66ae
diff --git a/tests/tests/telephony/current/src/android/telephony/cts/SubscriptionManagerTest.java b/tests/tests/telephony/current/src/android/telephony/cts/SubscriptionManagerTest.java
index 174af37..01401f8 100644
--- a/tests/tests/telephony/current/src/android/telephony/cts/SubscriptionManagerTest.java
+++ b/tests/tests/telephony/current/src/android/telephony/cts/SubscriptionManagerTest.java
@@ -757,13 +757,18 @@
                 // treat this as failure because it worked the first time
                 fail("setSubscriptionEnabled() did not work second time");
             }
+
+            // Reset default data subId as it may have been changed as part of the calls above
+            ShellIdentityUtils.invokeMethodWithShellPermissionsNoReturn(mSm,
+                    (sm) -> sm.setDefaultDataSubId(mSubId));
+
+            // Other tests also expect that cellular data must be available if telephony is
+            // supported. Wait for that before returning.
+            final CountDownLatch latch = waitForCellularNetwork();
+            latch.await(10, TimeUnit.SECONDS);
         } catch (InterruptedException e) {
             fail("InterruptedException");
         }
-
-        // Reset default data subId as it may have been changed as part of the calls above
-        ShellIdentityUtils.invokeMethodWithShellPermissionsNoReturn(mSm,
-                (sm) -> sm.setDefaultDataSubId(mSubId));
     }
 
     @Test
@@ -873,6 +878,28 @@
         return latch;
     }
 
+    /**
+     * Corresponding to findCellularNetwork()
+     */
+    private static CountDownLatch waitForCellularNetwork() {
+        final CountDownLatch latch = new CountDownLatch(1);
+        final ConnectivityManager cm = InstrumentationRegistry.getContext()
+                .getSystemService(ConnectivityManager.class);
+        cm.registerNetworkCallback(new NetworkRequest.Builder().build(),
+                new NetworkCallback() {
+                    @Override
+                    public void onCapabilitiesChanged(Network net, NetworkCapabilities caps) {
+                        if (caps.hasTransport(TRANSPORT_CELLULAR)
+                                && caps.hasCapability(NET_CAPABILITY_INTERNET)
+                                && caps.hasCapability(NET_CAPABILITY_NOT_RESTRICTED)) {
+                            latch.countDown();
+                            cm.unregisterNetworkCallback(this);
+                        }
+                    }
+                });
+        return latch;
+    }
+
     private static SubscriptionPlan buildValidSubscriptionPlan() {
         return SubscriptionPlan.Builder
                 .createRecurring(ZonedDateTime.parse("2007-03-14T00:00:00.000Z"),