1# Copyright 2018 The Chromium OS Authors. All rights reserved. 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5import socket 6import threading 7 8# Importing from six to maintain compatibility with Python 2. Safe to 9# `import queue` once Tauto transitions fully to Python 3. 10from six.moves import queue 11 12_BUF_SIZE = 4096 13 14class FakePrinter(): 15 """ 16 A fake printer (server). 17 18 It starts a thread that listens on given localhost's port and saves 19 incoming documents in the internal queue. Documents can be fetched from 20 the queue by calling the fetch_document() method. At the end, the printer 21 must be stopped by calling the stop() method. The stop() method is called 22 automatically when the object is managed by "with" statement. 23 See test_fake_printer.py for examples. 24 25 """ 26 27 def __init__(self, port): 28 """ 29 Initialize fake printer. 30 31 It configures the socket and starts the printer. If no exceptions 32 are thrown (the method succeeded), the printer must be stopped by 33 calling the stop() method. 34 35 @param port: port number on which the printer is supposed to listen 36 37 @raises socket or thread related exception in case of failure 38 39 """ 40 # If set to True, the printer is stopped either by invoking stop() 41 # method or by an internal error 42 self._stopped = False 43 # It is set when printer is stopped because of some internal error 44 self._error_message = None 45 # An internal queue with printed documents 46 self._documents = queue.Queue() 47 # Create a TCP/IP socket 48 self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 49 try: 50 # Bind the socket to the port 51 self._socket.bind( ('localhost', port) ) 52 # Start thread 53 self._thread = threading.Thread(target = self._thread_read_docs) 54 self._thread.start(); 55 except: 56 # failure - the socket must be closed before exit 57 self._socket.close() 58 raise 59 60 61 # These methods allow to use the 'with' statement to automaticaly stop 62 # the printer 63 def __enter__(self): 64 return self 65 def __exit__(self, exc_type, exc_value, traceback): 66 self.stop() 67 68 69 def stop(self): 70 """ 71 Stops the printer. 72 73 """ 74 self._stopped = True 75 self._thread.join() 76 77 78 def fetch_document(self, timeout): 79 """ 80 Fetches the next document from the internal queue. 81 82 This method returns the next document and removes it from the internal 83 queue. If there is no documents in the queue, it blocks until one 84 arrives. If waiting time exceeds a given timeout, an exception is 85 raised. 86 87 @param timeout: max waiting time in seconds 88 89 @returns next document from the internal queue 90 91 @raises Exception if the timeout was reached 92 93 """ 94 try: 95 return self._documents.get(block=True, timeout=timeout) 96 except queue.Empty: 97 # Builds a message for the exception 98 message = 'Timeout occured when waiting for the document. ' 99 if self._stopped: 100 message += 'The fake printer was stopped ' 101 if self._error_message is None: 102 message += 'by the stop() method.' 103 else: 104 message += 'because of the error: %s.' % self._error_message 105 else: 106 message += 'The fake printer is in valid state.' 107 # Raises and exception 108 raise Exception(message) 109 110 111 def _read_whole_document(self): 112 """ 113 Reads a document from the printer's socket. 114 115 It assumes that operation on sockets may timeout. 116 117 @returns whole document or None, if the printer was stopped 118 119 """ 120 # Accepts incoming connection 121 while True: 122 try: 123 (connection, client_address) = self._socket.accept() 124 # success - exit the loop 125 break 126 except socket.timeout: 127 # exit if the printer was stopped, else return to the loop 128 if self._stopped: 129 return None 130 131 # Reads document 132 document = bytearray() 133 while True: 134 try: 135 data = connection.recv(_BUF_SIZE) 136 # success - check data and continue 137 if not data: 138 # we got the whole document - exit the loop 139 break 140 # save chunk of the document and return to the loop 141 document.extend(data) 142 except socket.timeout: 143 # exit if the printer was stopped, else return to the loop 144 if self._stopped: 145 connection.close() 146 return None 147 148 # Closes connection & returns document 149 connection.close() 150 return bytes(document) 151 152 153 def _thread_read_docs(self): 154 """ 155 Reads documents from the printer's socket and adds them to the 156 internal queue. 157 158 It exits when the printer is stopped by the stop() method. 159 In case of any error (exception) it stops the printer and exits. 160 161 """ 162 try: 163 # Listen for incoming printer request. 164 self._socket.listen(1) 165 # All following socket's methods throw socket.timeout after 166 # 500 miliseconds 167 self._socket.settimeout(0.5) 168 169 while True: 170 # Reads document from the socket 171 document = self._read_whole_document() 172 # 'None' means that the printer was stopped -> exit 173 if document is None: 174 break 175 # Adds documents to the internal queue 176 self._documents.put(document) 177 except BaseException as e: 178 # Error occured, the printer must be stopped -> exit 179 self._error_message = str(e) 180 self._stopped = True 181 182 # Closes socket before the exit 183 self._socket.close() 184