Permalink
Browse files

add expect command

  • Loading branch information...
1 parent 2ec834a commit 49ff8c23bf31b5cfd8afa314f656469f29c8f524 Marc Sibson committed Jan 25, 2011
Showing with 93 additions and 18 deletions.
  1. +8 −0 tests/test_client.py
  2. +59 −12 vncdotool/client.py
  3. +26 −6 vncdotool/command.py
View
@@ -17,3 +17,11 @@ def test_capture_no_pil(self):
def test_multiple_captures(self):
self.client.capture('foo.png')
self.client.capture('bar.png')
+
+ def test_expect_initial_match(self):
+ self.client.expect('bar.png')
+
+ def test_expect_blocks_until_match(self):
+ pass
+ self.client.expect('bar.png')
+ # thousands of misses
View
@@ -10,6 +10,9 @@
from twisted.internet import reactor, defer
from twisted.internet.defer import Deferred
+import math
+import operator
+
KEYMAP = {
'bsp': rfb.KEY_BackSpace,
'tab': rfb.KEY_Tab,
@@ -89,6 +92,8 @@
class VNCDoToolClient(rfb.RFBClient):
x = 0
y = 0
+ screen = None
+ image = None
def keyPress(self, key):
""" Send a key press to the server
@@ -127,19 +132,47 @@ def captureScreen(self, filename):
# request initial screen update
self.framebufferUpdateRequest()
- self.capture = defer.Deferred()
- self.capture.addCallback(self._captureSave, filename)
+ self.deferred = defer.Deferred()
+ self.deferred.addCallback(self._captureSave, filename)
- return self.capture
+ return self.deferred
def _captureSave(self, data, filename):
- # lazy PIL import, allows other commands to work even if PIL
- # isn't installed
+ self.screen.save(filename)
+
+ return self
+
+ def expectScreen(self, filename, maxrms=0):
+ """ Wait until the display matches a target image
+
+ filename: an image file to read and compare against
+ maxrms: the maximum root mean square between histograms of the
+ screen and target image
+ """
+ self.framebufferUpdateRequest()
+
+ # lazy PIL, avoids breaking others
from PIL import Image
+ self.expected = Image.open(filename).histogram()
- size = (self.width, self.height)
- image = Image.fromstring('RGB', size, data, 'raw', 'RGBX')
- image.save(filename)
+ self.deferred = defer.Deferred()
+ self.deferred.addCallback(self._expectCompare, maxrms)
+
+ return self.deferred
+
+ def _expectCompare(self, data, maxrms):
+ hist = self.screen.histogram()
+
+ rms = math.sqrt(
+ reduce(operator.add,
+ map(lambda a,b: (a-b)**2, hist, self.expected)) / len(hist))
+
+ self.log('rms %d', rms)
+
+ if rms > maxrms:
+ self.deferred = defer.Deferred()
+ self.deferred.addCallback(self._expectCompare, maxrms)
+ return self.deferred
return self
@@ -172,17 +205,31 @@ def paste(self, message):
return self
def updateRectangle(self, x, y, width, height, data):
- assert (width, height) == (self.width, self.height)
- assert (x, y) == (0, 0)
+ # lazy PIL import, allows other commands to work even if PIL
+ # isn't installed
+ if not self.image:
+ from PIL import Image
+ self.image = Image
+
+
+ size = (width, height)
+ update = self.image.fromstring('RGB', size, data, 'raw', 'RGBX')
+ if not self.screen:
+ self.screen = update
+ else:
+ self.screen.paste(update, (x, y))
- self.capture.callback(data)
- self.capture = None
+ if self.deferred:
+ d = self.deferred
+ self.deferred = None
+ d.callback(data)
class VNCDoToolFactory(rfb.RFBFactory):
protocol = VNCDoToolClient
password = None
shared = 1
+ logger = None
def __init__(self):
self.deferred = Deferred()
View
@@ -11,17 +11,27 @@
from twisted.internet import reactor
from vncdotool.client import VNCDoToolFactory, VNCDoToolClient
+import sys
+
+def logger(fmt, *args):
+ print fmt % args
def log_connected(pcol):
print 'connected to', pcol.name
return pcol
+def error(reason):
+ try:
+ reason = reason.getErrorMessage()
+ except AttributeError:
+ pass
-def log_failed(reason):
- print 'connection failed', reason.getErrorMesssage()
-
+ print reason
+ reactor.exit_status = 10
+ reactor.stop()
def stop(pcol):
+ reactor.exit_status = 0
pcol.transport.loseConnection()
# XXX delay
reactor.callLater(0.1, reactor.stop)
@@ -52,10 +62,10 @@ def main():
f = VNCDoToolFactory()
if opts.verbose:
print 'connecting to %s:%d' % (opts.host, opts.port)
+ f.logger = logger
if opts.verbose:
f.deferred.addCallbacks(log_connected)
- f.deferred.addErrback(log_failed)
while args:
cmd = args.pop(0)
@@ -73,14 +83,24 @@ def main():
f.deferred.addCallback(VNCDoToolClient.keyPress, key)
elif cmd == 'capture':
filename = args.pop(0)
- f.deferred.addCallback(VNCDoToolClient.capture, filename)
+ f.deferred.addCallback(VNCDoToolClient.captureScreen, filename)
+ elif cmd == 'expect':
+ filename = args.pop(0)
+ rms = int(args.pop(0))
+ f.deferred.addCallback(VNCDoToolClient.expectScreen, filename, rms)
else:
print 'unknown cmd "%s"' % cmd
f.deferred.addCallback(stop)
+ f.deferred.addErrback(error)
+
+ d = reactor.connectTCP(opts.host, opts.port, f)
+
+ reactor.exit_status = 1
- reactor.connectTCP(opts.host, opts.port, f)
reactor.run()
+ sys.exit(reactor.exit_status)
+
if __name__ == '__main__':
main()

0 comments on commit 49ff8c2

Please sign in to comment.