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

   1 ##############################################################################
   2 # Copyright (c) 2016 ZTE Corporation
   3 # feng.xiaowei@zte.com.cn
   4 # All rights reserved. This program and the accompanying materials
   5 # are made available under the terms of the Apache License, Version 2.0
   6 # which accompanies this distribution, and is available at
   7 # http://www.apache.org/licenses/LICENSE-2.0
   8 ##############################################################################
   9 from bson.objectid import ObjectId
  10 from concurrent.futures import ThreadPoolExecutor
  11 from operator import itemgetter
  12
  13
  14 def thread_execute(method, *args, **kwargs):
  15         with ThreadPoolExecutor(max_workers=2) as executor:
  16             result = executor.submit(method, *args, **kwargs)
  17         return result
  18
  19
  20 class MemCursor(object):
  21     def __init__(self, collection):
  22         self.collection = collection
  23         self.count = len(self.collection)
  24         self.sorted = []
  25
  26     def _is_next_exist(self):
  27         return self.count != 0
  28
  29     @property
  30     def fetch_next(self):
  31         return thread_execute(self._is_next_exist)
  32
  33     def next_object(self):
  34         self.count -= 1
  35         return self.collection.pop()
  36
  37     def sort(self, key_or_list):
  38         key = key_or_list[0][0]
  39         if key_or_list[0][1] == -1:
  40             reverse = True
  41         else:
  42             reverse = False
  43
  44         if key_or_list is not None:
  45             self.collection = sorted(self.collection,
  46                                      key=itemgetter(key), reverse=reverse)
  47         return self
  48
  49     def limit(self, limit):
  50         if limit != 0 and limit < len(self.collection):
  51             self.collection = self.collection[0:limit]
  52             self.count = limit
  53         return self
  54
  55
  56 class MemDb(object):
  57
  58     def __init__(self):
  59         self.contents = []
  60         pass
  61
  62     def _find_one(self, spec_or_id=None, *args):
  63         if spec_or_id is not None and not isinstance(spec_or_id, dict):
  64             spec_or_id = {"_id": spec_or_id}
  65         if '_id' in spec_or_id:
  66             spec_or_id['_id'] = str(spec_or_id['_id'])
  67         cursor = self._find(spec_or_id, *args)
  68         for result in cursor:
  69             return result
  70         return None
  71
  72     def find_one(self, spec_or_id=None, *args):
  73         return thread_execute(self._find_one, spec_or_id, *args)
  74
  75     def _insert(self, doc_or_docs, check_keys=True):
  76
  77         docs = doc_or_docs
  78         return_one = False
  79         if isinstance(docs, dict):
  80             return_one = True
  81             docs = [docs]
  82
  83         ids = []
  84         for doc in docs:
  85             if '_id' not in doc:
  86                 doc['_id'] = str(ObjectId())
  87             if not check_keys or not self._find_one(doc['_id']):
  88                 ids.append(doc['_id'])
  89                 self.contents.append(doc_or_docs)
  90
  91         if len(ids) == 0:
  92             return None
  93         if return_one:
  94             return ids[0]
  95         else:
  96             return ids
  97
  98     def insert(self, doc_or_docs, check_keys=True):
  99         return thread_execute(self._insert, doc_or_docs, check_keys)
 100
 101     @staticmethod
 102     def _compare_date(spec, value):
 103         for k, v in spec.iteritems():
 104             if k == '$gte' and value >= v:
 105                 return True
 106         return False
 107
 108     @staticmethod
 109     def _in(content, *args):
 110         for arg in args:
 111             for k, v in arg.iteritems():
 112                 if k == 'start_date':
 113                     if not MemDb._compare_date(v, content.get(k)):
 114                         return False
 115                 elif k == 'trust_indicator':
 116                     if float(content.get(k)) != float(v):
 117                         return False
 118                 elif content.get(k, None) != v:
 119                     return False
 120
 121         return True
 122
 123     def _find(self, *args):
 124         res = []
 125         for content in self.contents:
 126             if self._in(content, *args):
 127                 res.append(content)
 128
 129         return res
 130
 131     def find(self, *args):
 132         return MemCursor(self._find(*args))
 133
 134     def _update(self, spec, document):
 135         updated = False
 136         for index in range(len(self.contents)):
 137             content = self.contents[index]
 138             if self._in(content, spec):
 139                 for k, v in document.iteritems():
 140                     updated = True
 141                     content[k] = v
 142             self.contents[index] = content
 143         return updated
 144
 145     def update(self, spec, document):
 146         return thread_execute(self._update, spec, document)
 147
 148     def _remove(self, spec_or_id=None):
 149         if spec_or_id is None:
 150             self.contents = []
 151         if not isinstance(spec_or_id, dict):
 152             spec_or_id = {'_id': spec_or_id}
 153         for index in range(len(self.contents)):
 154             content = self.contents[index]
 155             if self._in(content, spec_or_id):
 156                 del self.contents[index]
 157                 return True
 158         return False
 159
 160     def remove(self, spec_or_id=None):
 161         return thread_execute(self._remove, spec_or_id)
 162
 163     def clear(self):
 164         self._remove()
 165
 166 pods = MemDb()
 167 projects = MemDb()
 168 testcases = MemDb()
 169 results = MemDb()