adb: Add a way to reconnect TCP transports

This change adds a reconnect handler that tracks all TCP transports that
were connected at some point, but became disconnected. It does so by
attempting to reconnect every 10s for up to a minute.

Bug: 74411879
Test: system/core/adb/test_adb.py
Test: adb connect chromebook:22  # This runs with sslh
Test: CtsBootStatsTestCases
Test: emulator -show-kernel ; adb -s emulator-5554 shell

Change-Id: I7b9f6d181b71ccf5c26ff96c45d36aaf6409b992

1 files changed, 68 insertions, 16 deletions

diff --git a/test_adb.py b/test_adb.py
index 32bf0297..ce4d4ecf 100644
--- a/test_adb.py
+++ b/test_adb.py
@@ -75,9 +75,11 @@ def fake_adb_server(protocol=socket.AF_INET, port=0):
                 else:
                     # Client socket
                     data = r.recv(1024)
-                    if not data:
+                    if not data or data.startswith('OPEN'):
                         if r in cnxn_sent:
                             del cnxn_sent[r]
+                        r.shutdown(socket.SHUT_RDWR)
+                        r.close()
                         rlist.remove(r)
                         continue
                     if r in cnxn_sent:
@@ -97,6 +99,25 @@ def fake_adb_server(protocol=socket.AF_INET, port=0):
         server_thread.join()
 
 
+@contextlib.contextmanager
+def adb_connect(unittest, serial):
+    """Context manager for an ADB connection.
+
+    This automatically disconnects when done with the connection.
+    """
+
+    output = subprocess.check_output(['adb', 'connect', serial])
+    unittest.assertEqual(output.strip(), 'connected to {}'.format(serial))
+
+    try:
+        yield
+    finally:
+        # Perform best-effort disconnection. Discard the output.
+        p = subprocess.Popen(['adb', 'disconnect', serial],
+                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        p.communicate()
+
+
 class NonApiTest(unittest.TestCase):
     """Tests for ADB that aren't a part of the AndroidDevice API."""
 
@@ -278,29 +299,60 @@ class NonApiTest(unittest.TestCase):
         for protocol in (socket.AF_INET, socket.AF_INET6):
             try:
                 with fake_adb_server(protocol=protocol) as port:
-                    output = subprocess.check_output(
-                        ['adb', 'connect', 'localhost:{}'.format(port)])
-
-                    self.assertEqual(
-                        output.strip(), 'connected to localhost:{}'.format(port))
+                    serial = 'localhost:{}'.format(port)
+                    with adb_connect(self, serial):
+                        pass
             except socket.error:
                 print("IPv6 not available, skipping")
                 continue
 
     def test_already_connected(self):
-        with fake_adb_server() as port:
-            output = subprocess.check_output(
-                ['adb', 'connect', 'localhost:{}'.format(port)])
+        """Ensure that an already-connected device stays connected."""
 
-            self.assertEqual(
-                output.strip(), 'connected to localhost:{}'.format(port))
+        with fake_adb_server() as port:
+            serial = 'localhost:{}'.format(port)
+            with adb_connect(self, serial):
+                # b/31250450: this always returns 0 but probably shouldn't.
+                output = subprocess.check_output(['adb', 'connect', serial])
+                self.assertEqual(
+                    output.strip(), 'already connected to {}'.format(serial))
 
-            # b/31250450: this always returns 0 but probably shouldn't.
-            output = subprocess.check_output(
-                ['adb', 'connect', 'localhost:{}'.format(port)])
+    def test_reconnect(self):
+        """Ensure that a disconnected device reconnects."""
 
-            self.assertEqual(
-                output.strip(), 'already connected to localhost:{}'.format(port))
+        with fake_adb_server() as port:
+            serial = 'localhost:{}'.format(port)
+            with adb_connect(self, serial):
+                output = subprocess.check_output(['adb', '-s', serial,
+                                                  'get-state'])
+                self.assertEqual(output.strip(), 'device')
+
+                # This will fail.
+                p = subprocess.Popen(['adb', '-s', serial, 'shell', 'true'],
+                                     stdout=subprocess.PIPE,
+                                     stderr=subprocess.STDOUT)
+                output, _ = p.communicate()
+                self.assertEqual(output.strip(), 'error: closed')
+
+                subprocess.check_call(['adb', '-s', serial, 'wait-for-device'])
+
+                output = subprocess.check_output(['adb', '-s', serial,
+                                                  'get-state'])
+                self.assertEqual(output.strip(), 'device')
+
+                # Once we explicitly kick a device, it won't attempt to
+                # reconnect.
+                output = subprocess.check_output(['adb', 'disconnect', serial])
+                self.assertEqual(
+                    output.strip(), 'disconnected {}'.format(serial))
+                try:
+                    subprocess.check_output(['adb', '-s', serial, 'get-state'],
+                                            stderr=subprocess.STDOUT)
+                    self.fail('Device should not be available')
+                except subprocess.CalledProcessError as e:
+                    self.assertEqual(
+                        e.output.strip(),
+                        'error: device \'{}\' not found'.format(serial))
 
 def main():
     random.seed(0)