#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Version of WMCore/Services/Rucio intended to be used with mock or unittest.mock
"""
from __future__ import print_function, division
from future.utils import listitems
# from builtins import object # avoid importing this, it beraks things

import json
import logging
import os
import sys

from WMCore.Services.DBS.DBS3Reader import DBS3Reader, DBSReaderError
from WMCore.Services.Rucio.Rucio import WMRucioException, WMRucioDIDNotFoundException
from WMCore.WMBase import getTestBase
from WMQuality.Emulators.DataBlockGenerator.DataBlockGenerator import DataBlockGenerator

PROD_DBS = 'https://cmsweb-prod.cern.ch/dbs/prod/global/DBSReader'

NOT_EXIST_DATASET = 'thisdoesntexist'
# PILEUP_DATASET = '/HighPileUp/Run2011A-v1/RAW'
PILEUP_DATASET = '/GammaGammaToEE_Elastic_Pt15_8TeV-lpair/Summer12-START53_V7C-v1/GEN-SIM'

SITES = ['T2_XX_SiteA', 'T2_XX_SiteB', 'T2_XX_SiteC']
_BLOCK_LOCATIONS = {}

BLOCKS_PER_DATASET = 2
FILES_PER_BLOCK = 5
FILES_PER_DATASET = BLOCKS_PER_DATASET * FILES_PER_BLOCK

mockFile = os.path.join(getTestBase(), '..', 'data', 'Mock', 'RucioMockData.json')
try:
    with open(mockFile, 'r') as jsonObj:
        MOCK_DATA = json.load(jsonObj)
except IOError:
    MOCK_DATA = {}


def _unicode(data):
    """
    Temporary workaround for problem with how unicode strings are represented
    by unicode (as u"hello worls") and future.types.newstr (as "hello world")
    https://github.com/dmwm/WMCore/pull/10299#issuecomment-781600773
    """
    if sys.version_info[0] == 2:
        return unicode(data)
    else:
        return str(data)


class MockRucioApi(object):
    """
    Version of WMCore/Services/Rucio intended to be used with mock or unittest.mock
    """

    def __init__(self, acct, hostUrl=None, authUrl=None, configDict=None):
        print("Using MockRucioApi: acct={}, url={}, authUrl={}".format(acct, hostUrl, authUrl))
        configDict = configDict or {}
        self.dataBlocks = DataBlockGenerator()
        self.subRequests = {}

    def sitesByBlock(self, block):
        """
        Centralize the algorithm to decide where a block is based on the hash name
        :param block: the name of the block
        :return: a fake list of sites where the data is
        """
        logging.info("%s: Calling mock sitesByBlock", self.__class__.__name__)
        if hash(block) % 3 == 0:
            sites = ['T2_XX_SiteA']
        elif hash(block) % 3 == 1:
            sites = ['T2_XX_SiteA', 'T2_XX_SiteB']
        else:
            sites = ['T2_XX_SiteA', 'T2_XX_SiteB', 'T2_XX_SiteC']
        return sites

    def __getattr__(self, item):
        """
        __getattr__ gets called in case lookup of the actual method fails. We use this to return data based on
        a lookup table

        :param item: The method name the user is trying to call
        :return: The generic lookup function
        """

        def genericLookup(*args, **kwargs):
            """
            This function returns the mocked DBS data

            :param args: positional arguments it was called with
            :param kwargs: named arguments it was called with
            :return: the dictionary that DBS would have returned
            """
            logging.info("%s: Calling mock genericLookup", self.__class__.__name__)
            for key, value in listitems(kwargs):
                # json dumps/loads converts strings to unicode strings, do the same with kwargs
                if isinstance(value, str):
                    kwargs[key] = _unicode(value)
            if kwargs:
                signature = '%s:%s' % (item, sorted(kwargs.items()))
            else:
                signature = item

            try:
                return MOCK_DATA[signature]
            except KeyError:
                msg = "Rucio mock API failed to find key for signature: {}".format(signature)
                raise KeyError(msg)

        return genericLookup

    def getDataLockedAndAvailable(self, **kwargs):
        """
        Mock the method to discover where data is locked and available.
        Note that, by default, it will not return any Tape RSEs.
        :return: a unique list of RSEs
        """
        logging.info("%s: Calling mock getDataLockedAndAvailable", self.__class__.__name__)
        if 'name' not in kwargs:
            raise WMRucioException("A DID name must be provided to the getBlockLockedAndAvailable API")
        if self.isContainer(kwargs['name']):
            # then resolve it at container level and all its blocks
            return self.getContainerLockedAndAvailable(**kwargs)

        if 'grouping' in kwargs:
            # long strings seem not to be working, like ALL / DATASET. Make it short!
            kwargs['grouping'] = kwargs['grouping'][0]
        # TODO: either grouping or returnTape should change this response...
        returnTape = kwargs.pop("returnTape", False)

        rses = set()
        if kwargs['name'].split('#')[0] == PILEUP_DATASET:
            # Pileup is at a single site
            sites = ['T2_XX_SiteC']
            _BLOCK_LOCATIONS[kwargs['name']] = sites
        else:
            sites = self.sitesByBlock(block=kwargs['name'])
            _BLOCK_LOCATIONS[kwargs['name']] = sites
        rses.update(sites)
        return list(rses)

    def getPileupLockedAndAvailable(self, container, account, scope="cms"):
        """
        Mock method to resolve where the pileup container (and all its blocks)
        is locked and available.
        """
        logging.info("%s: calling mock getPileupLockedAndAvailable", self.__class__.__name__)
        result = dict()
        if not self.isContainer(container):
            raise WMRucioException("Pileup location needs to be resolved for a container DID type")

        kwargs = dict(name=container, account=account, scope=scope)

        try:
            DBS3Reader(PROD_DBS).checkDatasetPath(kwargs['name'])
            blocks = DBS3Reader(PROD_DBS).listFileBlocks(dataset=kwargs['name'])
            for block in blocks:
                result[block] = self.sitesByBlock(block)
        except DBSReaderError:
            logging.error("%s: Failed to fetch blocks from DBS", self.__class__.__name__)
        return result

    def getContainerLockedAndAvailable(self, **kwargs):
        """
        Mock the method to discover where container data is locked and available.
        Note that, by default, it will not return any Tape RSEs.
        :return: a unique list of RSEs
        """
        logging.info("%s: Calling mock getContainerLockedAndAvailable", self.__class__.__name__)
        if 'name' not in kwargs:
            raise WMRucioException("A DID name must be provided to the getContainerLockedAndAvailable API")
        kwargs.setdefault("scope", "cms")

        if kwargs['name'] == PILEUP_DATASET:
            return ['T2_XX_SiteA', 'T2_XX_SiteB', 'T2_XX_SiteC']
        try:
            DBS3Reader(PROD_DBS).checkDatasetPath(kwargs['name'])
            blocks = DBS3Reader(PROD_DBS).dbs.listBlocks(dataset=kwargs['name'])
            singleBlock = blocks[0]['block_name']
            return self.sitesByBlock(singleBlock)
        except DBSReaderError:
            return []

    def isContainer(self, didName, scope='cms'):
        """
        Mock check for whether a DID name corresponds to a container type or not,
        by simply relying on the naming convention
        :param didName: string with the DID name
        :param scope: string containing the Rucio scope (defaults to 'cms')
        :return: True if the DID is a container, else False
        """
        # TODO: figure use cases where we need to raise this exception
        if didName == "a bad DID name yet to be defined":
            msg = "Data identifier not found in MockRucio: {}".format(didName)
            raise WMRucioDIDNotFoundException(msg)
        return "#" not in didName

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Version of WMCore/Services/Rucio intended to be used with mock or unittest.mock
"""
from __future__ import print_function, division
from future.utils import listitems
# from builtins import object # avoid importing this, it beraks things

import json
import logging
import os
import sys

from WMCore.Services.DBS.DBS3Reader import DBS3Reader, DBSReaderError
from WMCore.Services.Rucio.Rucio import WMRucioException, WMRucioDIDNotFoundException
from WMCore.WMBase import getTestBase
from WMQuality.Emulators.DataBlockGenerator.DataBlockGenerator import DataBlockGenerator

PROD_DBS = 'https://cmsweb-prod.cern.ch/dbs/prod/global/DBSReader'

NOT_EXIST_DATASET = 'thisdoesntexist'
# PILEUP_DATASET = '/HighPileUp/Run2011A-v1/RAW'
PILEUP_DATASET = '/GammaGammaToEE_Elastic_Pt15_8TeV-lpair/Summer12-START53_V7C-v1/GEN-SIM'

SITES = ['T2_XX_SiteA', 'T2_XX_SiteB', 'T2_XX_SiteC']
_BLOCK_LOCATIONS = {}

BLOCKS_PER_DATASET = 2
FILES_PER_BLOCK = 5
FILES_PER_DATASET = BLOCKS_PER_DATASET * FILES_PER_BLOCK

mockFile = os.path.join(getTestBase(), '..', 'data', 'Mock', 'RucioMockData.json')
try:
    with open(mockFile, 'r') as jsonObj:
        MOCK_DATA = json.load(jsonObj)
except IOError:
    MOCK_DATA = {}


def _unicode(data):
    """
    Temporary workaround for problem with how unicode strings are represented
    by unicode (as u"hello worls") and future.types.newstr (as "hello world")
    https://github.com/dmwm/WMCore/pull/10299#issuecomment-781600773
    """
    if sys.version_info[0] == 2:
        return unicode(data)
    else:
        return str(data)


class MockRucioApi(object):
    """
    Version of WMCore/Services/Rucio intended to be used with mock or unittest.mock
    """

    def __init__(self, acct, hostUrl=None, authUrl=None, configDict=None):
        print("Using MockRucioApi: acct={}, url={}, authUrl={}".format(acct, hostUrl, authUrl))
        configDict = configDict or {}
        self.dataBlocks = DataBlockGenerator()
        self.subRequests = {}

    def sitesByBlock(self, block):
        """
        Centralize the algorithm to decide where a block is based on the hash name
        :param block: the name of the block
        :return: a fake list of sites where the data is
        """
        logging.info("%s: Calling mock sitesByBlock", self.__class__.__name__)
        if hash(block) % 3 == 0:
            sites = ['T2_XX_SiteA']
        elif hash(block) % 3 == 1:
            sites = ['T2_XX_SiteA', 'T2_XX_SiteB']
        else:
            sites = ['T2_XX_SiteA', 'T2_XX_SiteB', 'T2_XX_SiteC']
        return sites

    def __getattr__(self, item):
        """
        __getattr__ gets called in case lookup of the actual method fails. We use this to return data based on
        a lookup table

        :param item: The method name the user is trying to call
        :return: The generic lookup function
        """

        def genericLookup(*args, **kwargs):
            """
            This function returns the mocked DBS data

            :param args: positional arguments it was called with
            :param kwargs: named arguments it was called with
            :return: the dictionary that DBS would have returned
            """
            logging.info("%s: Calling mock genericLookup", self.__class__.__name__)
            for key, value in listitems(kwargs):
                # json dumps/loads converts strings to unicode strings, do the same with kwargs
                if isinstance(value, str):
                    kwargs[key] = _unicode(value)
            if kwargs:
                signature = '%s:%s' % (item, sorted(kwargs.items()))
            else:
                signature = item

            try:
                return MOCK_DATA[signature]
            except KeyError:
                msg = "Rucio mock API failed to find key for signature: {}".format(signature)
                raise KeyError(msg)

        return genericLookup

    def getDataLockedAndAvailable(self, **kwargs):
        """
        Mock the method to discover where data is locked and available.
        Note that, by default, it will not return any Tape RSEs.
        :return: a unique list of RSEs
        """
        logging.info("%s: Calling mock getDataLockedAndAvailable", self.__class__.__name__)
        if 'name' not in kwargs:
            raise WMRucioException("A DID name must be provided to the getBlockLockedAndAvailable API")
        if self.isContainer(kwargs['name']):
            # then resolve it at container level and all its blocks
            return self.getContainerLockedAndAvailable(**kwargs)

        if 'grouping' in kwargs:
            # long strings seem not to be working, like ALL / DATASET. Make it short!
            kwargs['grouping'] = kwargs['grouping'][0]
        # TODO: either grouping or returnTape should change this response...
        returnTape = kwargs.pop("returnTape", False)

        rses = set()
        if kwargs['name'].split('#')[0] == PILEUP_DATASET:
            # Pileup is at a single site
            sites = ['T2_XX_SiteC']
            _BLOCK_LOCATIONS[kwargs['name']] = sites
        else:
            sites = self.sitesByBlock(block=kwargs['name'])
            _BLOCK_LOCATIONS[kwargs['name']] = sites
        rses.update(sites)
        return list(rses)

    def getPileupLockedAndAvailable(self, container, account, scope="cms"):
        """
        Mock method to resolve where the pileup container (and all its blocks)
        is locked and available.
        """
        logging.info("%s: calling mock getPileupLockedAndAvailable", self.__class__.__name__)
        result = dict()
        if not self.isContainer(container):
            raise WMRucioException("Pileup location needs to be resolved for a container DID type")

        kwargs = dict(name=container, account=account, scope=scope)

        try:
            DBS3Reader(PROD_DBS).checkDatasetPath(kwargs['name'])
            blocks = DBS3Reader(PROD_DBS).listFileBlocks(dataset=kwargs['name'])
            for block in blocks:
                result[block] = self.sitesByBlock(block)
        except DBSReaderError:
            logging.error("%s: Failed to fetch blocks from DBS", self.__class__.__name__)
        return result

    def getContainerLockedAndAvailable(self, **kwargs):
        """
        Mock the method to discover where container data is locked and available.
        Note that, by default, it will not return any Tape RSEs.
        :return: a unique list of RSEs
        """
        logging.info("%s: Calling mock getContainerLockedAndAvailable", self.__class__.__name__)
        if 'name' not in kwargs:
            raise WMRucioException("A DID name must be provided to the getContainerLockedAndAvailable API")
        kwargs.setdefault("scope", "cms")

        if kwargs['name'] == PILEUP_DATASET:
            return ['T2_XX_SiteA', 'T2_XX_SiteB', 'T2_XX_SiteC']
        try:
            DBS3Reader(PROD_DBS).checkDatasetPath(kwargs['name'])
            blocks = DBS3Reader(PROD_DBS).dbs.listBlocks(dataset=kwargs['name'])
            singleBlock = blocks[0]['block_name']
            return self.sitesByBlock(singleBlock)
        except DBSReaderError:
            return []

    def isContainer(self, didName, scope='cms'):
        """
        Mock check for whether a DID name corresponds to a container type or not,
        by simply relying on the naming convention
        :param didName: string with the DID name
        :param scope: string containing the Rucio scope (defaults to 'cms')
        :return: True if the DID is a container, else False
        """
        # TODO: figure use cases where we need to raise this exception
        if didName == "a bad DID name yet to be defined":
            msg = "Data identifier not found in MockRucio: {}".format(didName)
            raise WMRucioDIDNotFoundException(msg)
        return "#" not in didName

#!/usr/bin/env python # -*- coding: utf-8 -*- Version of WMCore/Services/Rucio intended to be used with mock or unittest.mock # from builtins import object # avoid importing this, it beraks things # PILEUP_DATASET = '/HighPileUp/Run2011A-v1/RAW' Temporary workaround for problem with how unicode strings are represented     by unicode (as u"hello worls") and future.types.newstr (as "hello world")     https://github.com/dmwm/WMCore/pull/10299#issuecomment-781600773 Version of WMCore/Services/Rucio intended to be used with mock or unittest.mock Centralize the algorithm to decide where a block is based on the hash name         :param block: the name of the block         :return: a fake list of sites where the data is __getattr__ gets called in case lookup of the actual method fails. We use this to return data based on         a lookup table          :param item: The method name the user is trying to call         :return: The generic lookup function This function returns the mocked DBS data              :param args: positional arguments it was called with             :param kwargs: named arguments it was called with             :return: the dictionary that DBS would have returned # json dumps/loads converts strings to unicode strings, do the same with kwargs Mock the method to discover where data is locked and available.         Note that, by default, it will not return any Tape RSEs.         :return: a unique list of RSEs # then resolve it at container level and all its blocks # long strings seem not to be working, like ALL / DATASET. Make it short! # TODO: either grouping or returnTape should change this response... # Pileup is at a single site Mock method to resolve where the pileup container (and all its blocks)         is locked and available. Mock the method to discover where container data is locked and available.         Note that, by default, it will not return any Tape RSEs.         :return: a unique list of RSEs Mock check for whether a DID name corresponds to a container type or not,         by simply relying on the naming convention         :param didName: string with the DID name         :param scope: string containing the Rucio scope (defaults to 'cms')         :return: True if the DID is a container, else False # TODO: figure use cases where we need to raise this exception