utils/test/result_collection_api/tests/unit/fake_pymongo.py

   1 from bson.objectid import ObjectId
   2 from concurrent.futures import ThreadPoolExecutor
   3
   4
   5 __author__ = 'serena'
   6
   7
   8 def thread_execute(method, *args, **kwargs):
   9         with ThreadPoolExecutor(max_workers=2) as executor:
  10             result = executor.submit(method, *args, **kwargs)
  11         return result
  12
  13
  14 class MemCursor(object):
  15     def __init__(self, collection):
  16         self.collection = collection
  17         self.count = len(self.collection)
  18
  19     def _is_next_exist(self):
  20         return self.count != 0
  21
  22     @property
  23     def fetch_next(self):
  24         return thread_execute(self._is_next_exist)
  25
  26     def next_object(self):
  27         self.count -= 1
  28         return self.collection.pop()
  29
  30
  31 class MemDb(object):
  32
  33     def __init__(self):
  34         self.contents = []
  35         pass
  36
  37     def _find_one(self, spec_or_id=None, *args):
  38         if spec_or_id is not None and not isinstance(spec_or_id, dict):
  39             spec_or_id = {"_id": spec_or_id}
  40         if '_id' in spec_or_id:
  41             spec_or_id['_id'] = str(spec_or_id['_id'])
  42         cursor = self._find(spec_or_id, *args)
  43         for result in cursor:
  44             return result
  45         return None
  46
  47     def find_one(self, spec_or_id=None, *args):
  48         return thread_execute(self._find_one, spec_or_id, *args)
  49
  50     def _insert(self, doc_or_docs, check_keys=True):
  51
  52         docs = doc_or_docs
  53         return_one = False
  54         if isinstance(docs, dict):
  55             return_one = True
  56             docs = [docs]
  57
  58         ids = []
  59         for doc in docs:
  60             if '_id' not in doc:
  61                 doc['_id'] = str(ObjectId())
  62             if not check_keys or not self._find_one(doc['_id']):
  63                 ids.append(doc['_id'])
  64                 self.contents.append(doc_or_docs)
  65
  66         if len(ids) == 0:
  67             return None
  68         if return_one:
  69             return ids[0]
  70         else:
  71             return ids
  72
  73     def insert(self, doc_or_docs, check_keys=True):
  74         return thread_execute(self._insert, doc_or_docs, check_keys)
  75
  76     @staticmethod
  77     def _in(content, *args):
  78         for arg in args:
  79             for k, v in arg.iteritems():
  80                 if k != 'creation_date' and content.get(k, None) != v:
  81                     return False
  82
  83         return True
  84
  85     def _find(self, *args):
  86         res = []
  87         for content in self.contents:
  88             if self._in(content, *args):
  89                 res.append(content)
  90
  91         return res
  92
  93     def find(self, *args):
  94         return MemCursor(self._find(*args))
  95
  96     def _update(self, spec, document):
  97         updated = False
  98         for index in range(len(self.contents)):
  99             content = self.contents[index]
 100             if self._in(content, spec):
 101                 for k, v in document.iteritems():
 102                     updated = True
 103                     content[k] = v
 104             self.contents[index] = content
 105         return updated
 106
 107     def update(self, spec, document):
 108         return thread_execute(self._update, spec, document)
 109
 110     def _remove(self, spec_or_id=None):
 111         if spec_or_id is None:
 112             self.contents = []
 113         if not isinstance(spec_or_id, dict):
 114             spec_or_id = {'_id': spec_or_id}
 115         for index in range(len(self.contents)):
 116             content = self.contents[index]
 117             if self._in(content, spec_or_id):
 118                 del self.contents[index]
 119                 return True
 120         return False
 121
 122     def remove(self, spec_or_id=None):
 123         return thread_execute(self._remove, spec_or_id)
 124
 125     def clear(self):
 126         self._remove()
 127
 128 pods = MemDb()
 129 projects = MemDb()
 130 testcases = MemDb()
 131 results = MemDb()