From 15a77a841d924653b20b0ab05b776a5add1713f8 Mon Sep 17 00:00:00 2001
From: Marc Bonnici <marc.bonnici@arm.com>
Date: Thu, 5 Dec 2019 11:09:00 +0000
Subject: [PATCH] collector/screencapture: Refactor to use new collector
 interface

Update the interface to make use of the collector interface.
Notable changes are the removal of the `output_path` path provided on
initialisation which will now be provided by the dedicated `set_output`
method.
---
 devlib/collector/screencapture.py | 31 +++++++++++++++++++++++--------
 1 file changed, 23 insertions(+), 8 deletions(-)

diff --git a/devlib/collector/screencapture.py b/devlib/collector/screencapture.py
index 0a5ca64..399227f 100644
--- a/devlib/collector/screencapture.py
+++ b/devlib/collector/screencapture.py
@@ -19,13 +19,14 @@ import sys
 import threading
 import time
 
-from devlib.collector import BaseCollector
+from devlib.collector import (CollectorBase, CollectorOutput,
+                              CollectorOutputEntry)
 from devlib.exception import WorkerThreadError
 
 
 class ScreenCapturePoller(threading.Thread):
 
-    def __init__(self, target, period, output_path=None, timeout=30):
+    def __init__(self, target, period, timeout=30):
         super(ScreenCapturePoller, self).__init__()
         self.target = target
         self.logger = logging.getLogger('screencapture')
@@ -36,11 +37,16 @@ class ScreenCapturePoller(threading.Thread):
         self.last_poll = 0
         self.daemon = True
         self.exc = None
+        self.output_path = None
+
+    def set_output(self, output_path):
         self.output_path = output_path
 
     def run(self):
         self.logger.debug('Starting screen capture polling')
         try:
+            if self.output_path is None:
+                raise RuntimeError("Output path was not set.")
             while True:
                 if self.stop_signal.is_set():
                     break
@@ -66,24 +72,33 @@ class ScreenCapturePoller(threading.Thread):
         self.target.capture_screen(os.path.join(self.output_path, "screencap_{ts}.png"))
 
 
-class ScreenCaptureCollector(BaseCollector):
+class ScreenCaptureCollector(CollectorBase):
 
-    def __init__(self, target, output_path=None, period=None):
+    def __init__(self, target, period=None):
         super(ScreenCaptureCollector, self).__init__(target)
         self._collecting = False
-        self.output_path = output_path
+        self.output_path = None
         self.period = period
         self.target = target
-        self._poller = ScreenCapturePoller(self.target, self.period,
-                                           self.output_path)
+
+    def set_output(self, output_path):
+        self.output_path = output_path
 
     def reset(self):
-        pass
+        self._poller = ScreenCapturePoller(self.target, self.period)
+
+    def get_data(self):
+        if self.output_path is None:
+            raise RuntimeError("No data collected.")
+        return CollectorOutput([CollectorOutputEntry(self.output_path, 'directory')])
 
     def start(self):
         """
         Start collecting the screenshots
         """
+        if self.output_path is None:
+            raise RuntimeError("Output path was not set.")
+        self._poller.set_output(self.output_path)
         self._poller.start()
         self._collecting = True