/src/python/WMQuality/Emulators/DBSClient/MockDbsApi.py

https://github.com/dmwm/WMCore
Python | 143 lines | 95 code | 23 blank | 25 comment | 21 complexity | 6bdc7ef4fa4c4a05786ae79a8bbca66d MD5 | raw file
  1. #!/usr/bin/env python
  2. """
  3. Version of dbsClient.dbsApi intended to be used with mock or unittest.mock
  4. """
  5. from __future__ import (division, print_function)
  6. from builtins import object
  7. from future.utils import viewitems
  8. import copy
  9. import json
  10. import os
  11. from RestClient.ErrorHandling.RestClientExceptions import HTTPError
  12. from WMCore.Services.DBS.DBSErrors import DBSReaderError
  13. from WMCore.WMBase import getTestBase
  14. from Utils.Utilities import encodeUnicodeToBytesConditional
  15. from Utils.PythonVersion import PY2
  16. # Read in the data just once so that we don't have to do it for every test (in __init__)
  17. mockData = {}
  18. globalFile = os.path.join(getTestBase(), '..', 'data', 'Mock', 'DBSMockData.json')
  19. phys03File = os.path.join(getTestBase(), '..', 'data', 'Mock', 'DBSMockData03.json')
  20. try:
  21. with open(globalFile, 'r') as mockFile:
  22. mockDataGlobal = json.load(mockFile)
  23. except IOError:
  24. mockDataGlobal = {}
  25. try:
  26. with open(phys03File, 'r') as mockFile:
  27. mockData03 = json.load(mockFile)
  28. except IOError:
  29. mockData03 = {}
  30. mockData['https://cmsweb-prod.cern.ch/dbs/prod/global/DBSReader'] = mockDataGlobal
  31. mockData['https://cmsweb-prod.cern.ch/dbs/prod/phys03/DBSReader'] = mockData03
  32. class MockDbsApi(object):
  33. def __init__(self, url):
  34. print("Using MockDBSApi")
  35. self.url = url.strip('/')
  36. # print("Initializing MockDBSApi")
  37. def serverinfo(self):
  38. return {'dbs_instance': 'MOCK', 'dbs_version': '3.3.144'}
  39. def listFileArray(self, **kwargs):
  40. """
  41. Handle the case when logical_file_name is called with a list (longer than one) of files
  42. since we don't want to store all permutations. Rebuild the list of dicts that DBS returns
  43. Args:
  44. **kwargs: any kwargs that dbs client accepts
  45. Returns:
  46. """
  47. self.item = 'listFileArray'
  48. if 'logical_file_name' in kwargs and len(kwargs['logical_file_name']) > 1:
  49. origArgs = copy.deepcopy(kwargs)
  50. returnDicts = []
  51. for lfn in kwargs['logical_file_name']:
  52. origArgs.update({'logical_file_name': [lfn]})
  53. returnDicts.extend(self.genericLookup(**origArgs))
  54. return returnDicts
  55. else:
  56. return self.genericLookup(**kwargs)
  57. def listFileLumiArray(self, **kwargs):
  58. """
  59. Handle the case when logical_file_name is called with a list (longer than one) of files
  60. since we don't want to store all permutations. Rebuild the list of dicts that DBS returns
  61. Args:
  62. **kwargs: any kwargs that dbs client accepts
  63. Returns:
  64. """
  65. self.item = 'listFileLumiArray'
  66. if 'logical_file_name' in kwargs and len(kwargs['logical_file_name']) > 1:
  67. origArgs = copy.deepcopy(kwargs)
  68. returnDicts = []
  69. # since we iterate over this, we better make sure it's a list to avoid
  70. # things like: ['U', 'N', 'K', 'N', 'O', 'W', 'N']
  71. if isinstance(kwargs['logical_file_name'], str):
  72. kwargs['logical_file_name'] = [kwargs['logical_file_name']]
  73. for lfn in kwargs['logical_file_name']:
  74. origArgs.update({'logical_file_name': [lfn]})
  75. returnDicts.extend(self.genericLookup(**origArgs))
  76. return returnDicts
  77. else:
  78. return self.genericLookup(**kwargs)
  79. def __getattr__(self, item):
  80. """
  81. __getattr__ gets called in case lookup of the actual method fails. We use this to return data based on
  82. a lookup table
  83. :param item: The method name the user is trying to call
  84. :return: The generic lookup function
  85. """
  86. self.item = item
  87. return self.genericLookup
  88. def genericLookup(self, *args, **kwargs):
  89. """
  90. This function returns the mocked DBS data
  91. :param args: positional arguments it was called with
  92. :param kwargs: named arguments it was called with
  93. :return: the dictionary that DBS would have returned
  94. """
  95. if self.url not in mockData:
  96. raise DBSReaderError("Mock DBS emulator knows nothing about instance %s" % self.url)
  97. if kwargs:
  98. for k in kwargs:
  99. if isinstance(kwargs[k], (list, tuple)):
  100. kwargs[k] = [encodeUnicodeToBytesConditional(item, condition=PY2) for item in kwargs[k]]
  101. else:
  102. kwargs[k] = encodeUnicodeToBytesConditional(kwargs[k], condition=PY2)
  103. signature = '%s:%s' % (self.item, sorted(viewitems(kwargs)))
  104. else:
  105. signature = self.item
  106. try:
  107. if mockData[self.url][signature] == 'Raises HTTPError':
  108. raise HTTPError('http:/dbs.mock.fail', 400, 'MockDBS is raising an exception in place of DBS', 'Dummy header', 'Dummy body')
  109. else:
  110. return mockData[self.url][signature]
  111. except KeyError:
  112. if kwargs.get('dataset', None) == '/HighPileUp/Run2011A-v1/RAW-BLAH':
  113. return []
  114. raise KeyError("DBS mock API could not return data for method %s, args=%s, and kwargs=%s (URL %s) (Signature: %s)" %
  115. (self.item, args, kwargs, self.url, signature))