13804: Fix tests

[arvados.git] / tools / crunchstat-summary / crunchstat_summary / reader.py

diff --git a/tools/crunchstat-summary/crunchstat_summary/reader.py b/tools/crunchstat-summary/crunchstat_summary/reader.py
index 564e8208f6c03cc337d1c65b9990622aa3a7eaee..e8f0861be49350eddc232bd8826a4681af3a424f 100644 (file)
--- a/tools/crunchstat-summary/crunchstat_summary/reader.py
+++ b/tools/crunchstat-summary/crunchstat_summary/reader.py
@@ -1,30 +1,46 @@
+# Copyright (C) The Arvados Authors. All rights reserved.
+#
+# SPDX-License-Identifier: AGPL-3.0
+
 from __future__ import print_function
 
 import arvados
 import Queue
 import threading
+import _strptime
 
 from crunchstat_summary import logger
 
 
 class CollectionReader(object):
     def __init__(self, collection_id):
-        logger.debug('load collection %s', collection_id)
-        collection = arvados.collection.CollectionReader(collection_id)
-        filenames = [filename for filename in collection]
-        if len(filenames) != 1:
-            raise ValueError(
-                "collection {} has {} files; need exactly one".format(
-                    collection_id, len(filenames)))
-        self._reader = collection.open(filenames[0])
-        self._label = "{}/{}".format(collection_id, filenames[0])
+        self._collection_id = collection_id
+        self._label = collection_id
+        self._reader = None
 
     def __str__(self):
         return self._label
 
     def __iter__(self):
+        logger.debug('load collection %s', self._collection_id)
+        collection = arvados.collection.CollectionReader(self._collection_id)
+        filenames = [filename for filename in collection]
+        if len(filenames) == 1:
+            filename = filenames[0]
+        else:
+            filename = 'crunchstat.txt'
+        self._label = "{}/{}".format(self._collection_id, filename)
+        self._reader = collection.open(filename)
         return iter(self._reader)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self._reader:
+            self._reader.close()
+            self._reader = None
+
 
 class LiveLogReader(object):
     EOF = None
@@ -48,6 +64,7 @@ class LiveLogReader(object):
                     limit=1000,
                     order=['id asc'],
                     filters=filters + [['id','>',str(last_id)]],
+                    select=['id', 'properties'],
                 ).execute(num_retries=2)
                 got += len(page['items'])
                 logger.debug(
@@ -77,3 +94,9 @@ class LiveLogReader(object):
             self._thread.join()
             raise StopIteration
         return line
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass