michael@0: #!/usr/bin/env python
michael@0: 
michael@0: # Any copyright is dedicated to the Public Domain.
michael@0: # http://creativecommons.org/publicdomain/zero/1.0/
michael@0: 
michael@0: import datetime
michael@0: import socket
michael@0: import time
michael@0: 
michael@0: from threading import Thread
michael@0: 
michael@0: class MockAgent(object):
michael@0: 
michael@0:     MAX_WAIT_TIME_SECONDS = 10
michael@0:     SOCKET_TIMEOUT_SECONDS = 5
michael@0: 
michael@0:     def __init__(self, tester, start_commands = None, commands = []):
michael@0:         if start_commands:
michael@0:             self.commands = start_commands
michael@0:         else:
michael@0:             self.commands = [("testroot", "/mnt/sdcard"),
michael@0:                              ("isdir /mnt/sdcard/tests", "TRUE"),
michael@0:                              ("ver", "SUTAgentAndroid Version 1.14")]
michael@0:         self.commands = self.commands + commands
michael@0: 
michael@0:         self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
michael@0:         self._sock.bind(("127.0.0.1", 0))
michael@0:         self._sock.listen(1)
michael@0: 
michael@0:         self.tester = tester
michael@0: 
michael@0:         self.thread = Thread(target=self._serve_thread)
michael@0:         self.thread.start()
michael@0: 
michael@0:         self.should_stop = False
michael@0: 
michael@0:     @property
michael@0:     def port(self):
michael@0:         return self._sock.getsockname()[1]
michael@0: 
michael@0:     def _serve_thread(self):
michael@0:         conn = None
michael@0:         while self.commands:
michael@0:             if not conn:
michael@0:                 conn, addr = self._sock.accept()
michael@0:                 conn.settimeout(self.SOCKET_TIMEOUT_SECONDS)
michael@0:                 conn.send("$>\x00")
michael@0:             (command, response) = self.commands.pop(0)
michael@0:             data = ''
michael@0:             timeout = datetime.datetime.now() + datetime.timedelta(
michael@0:                 seconds=self.MAX_WAIT_TIME_SECONDS)
michael@0:             # The data might come in chunks, particularly if we are expecting
michael@0:             # multiple lines, as with push commands.
michael@0:             while (len(data) < len(command) and
michael@0:                    datetime.datetime.now() < timeout):
michael@0:                 try:
michael@0:                     data += conn.recv(1024)
michael@0:                 except socket.timeout:
michael@0:                     # We handle timeouts in the main loop.
michael@0:                     pass
michael@0:             self.tester.assertEqual(data.strip(), command)
michael@0:             # send response and prompt separately to test for bug 789496
michael@0:             # FIXME: Improve the mock agent, since overloading the meaning
michael@0:             # of 'response' is getting confusing.
michael@0:             if response is None: # code for "shut down"
michael@0:                 conn.shutdown(socket.SHUT_RDWR)
michael@0:                 conn.close()
michael@0:                 conn = None
michael@0:             elif type(response) is int: # code for "time out"
michael@0:                 max_timeout = 15.0
michael@0:                 timeout = 0.0
michael@0:                 interval = 0.1
michael@0:                 while not self.should_stop and timeout < max_timeout:
michael@0:                     time.sleep(interval)
michael@0:                     timeout += interval
michael@0:                 if timeout >= max_timeout:
michael@0:                     raise Exception("Maximum timeout reached! This should not "
michael@0:                                     "happen")
michael@0:                 return
michael@0:             else:
michael@0:                 # pull is handled specially, as we just pass back the full
michael@0:                 # command line
michael@0:                 if "pull" in command:
michael@0:                     conn.send(response)
michael@0:                 else:
michael@0:                     conn.send("%s\n" % response)
michael@0:                 conn.send("$>\x00")
michael@0: 
michael@0:     def wait(self):
michael@0:         self.thread.join()