###############################################################################
# (c) Copyright 2019 CERN for the benefit of the LHCb Collaboration           #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "LICENSE".   #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################
"""Unit tests for Workflow Modules."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

__RCSID__ = "$Id$"

# pylint: disable=protected-access, missing-docstring, invalid-name, line-too-long

import unittest
import os
import copy
import shutil

from mock import MagicMock, patch

from DIRAC.DataManagementSystem.Client.test.mock_DM import dm_mock
from DIRAC.Resources.Catalog.test.mock_FC import fc_mock

from DIRAC import gLogger
from DIRAC.RequestManagementSystem.Client.Request import Request

# mocks
from LHCbDIRAC.Workflow.Modules.mock_Commons import prod_id, prod_job_id, wms_job_id, \
    workflowStatus, stepStatus, step_id, step_number,\
    step_commons, wf_commons,\
    rc_mock
from LHCbDIRAC.BookkeepingSystem.Client.test.mock_BookkeepingClient import bkc_mock

# sut
from LHCbDIRAC.Workflow.Modules.FailoverRequest import FailoverRequest
from LHCbDIRAC.Workflow.Modules.RemoveInputData import RemoveInputData
from LHCbDIRAC.Workflow.Modules.UserJobFinalization import UserJobFinalization
from LHCbDIRAC.Workflow.Modules.StepAccounting import StepAccounting
from LHCbDIRAC.Workflow.Modules.UploadLogFile import UploadLogFile
from LHCbDIRAC.Workflow.Modules.FileUsage import FileUsage
from LHCbDIRAC.Workflow.Modules.CreateDataFile import CreateDataFile

getDestinationSEListMock = MagicMock()
getDestinationSEListMock.return_value = []
getDestinationSEListMockCNAF = MagicMock()
getDestinationSEListMockCNAF.return_value = ['CNAF']


class ModulesTestCase(unittest.TestCase):
  """Base class for the Modules test cases."""

  def setUp(self):

    gLogger.setLevel('DEBUG')
    self.maxDiff = None

    self.jsu_mock = MagicMock()
    self.jsu_mock.setJobApplicationStatus.return_value = {'OK': True, 'Value': ''}

    self.jsu_mock = MagicMock()
    self.jsu_mock.setJobApplicationStatus.return_value = {'OK': True, 'Value': ''}

    self.ft_mock = MagicMock()
    self.ft_mock.transferAndRegisterFile.return_value = {'OK': True, 'Value': {'uploadedSE': ''}}
    self.ft_mock.transferAndRegisterFileFailover.return_value = {'OK': True, 'Value': {}}
    self.ft_mock.request = rc_mock
    self.ft_mock.FileCatalog = fc_mock

    self.nc_mock = MagicMock()
    self.nc_mock.sendMail.return_value = {'OK': True, 'Value': ''}

    self.xf_o_mock = MagicMock()
    self.xf_o_mock.inputFileStats = {'a': 1, 'b': 2}
    self.xf_o_mock.outputFileStats = {'a': 1, 'b': 2}
    self.xf_o_mock.analyse.return_value = True

    self.jobStep_mock = MagicMock()
    self.jobStep_mock.commit.return_value = {'OK': True, 'Value': ''}
    self.jobStep_mock.setValuesFromDict.return_value = {'OK': True, 'Value': ''}
    self.jobStep_mock.checkValues.return_value = {'OK': True, 'Value': ''}

  def tearDown(self):
    for fileProd in ['appLog',
                     'foo.txt',
                     'aaa.Bhadron.dst',
                     'bbb.Calibration.dst',
                     'bar.py',
                     'aLongLog.log',
                     'bookkeeping_123_00000456_321.xml',
                     'aLongLog.log.gz',
                     'ccc.charm.mdst',
                     'ccc.charm.mdst',
                     'prova.txt',
                     'aLog.log',
                     'BAR.txt',
                     'FooBAR.ext.txt',
                     'foo_1.txt',
                     'bar_2.py',
                     'bar.txt',
                     'ErrorLogging_Step1_coredump.log',
                     '123_00000456_request.xml',
                     'lfn1',
                     'lfn2',
                     'XMLSummaryFile',
                     'aaa.bhadron.dst',
                     'bbb.calibration.dst',
                     'ProductionOutputData',
                     'data.py',
                     '123_00000456_request.json',
                     '00000123_00000456.tar',
                     'someOtherDir',
                     'DISABLE_WATCHDOG_CPU_WALLCLOCK_CHECK',
                     'myfoo.blah',
                     'prodConf_someApp__.py',
                     'prodConf_someApp___.py']:
      try:
        os.remove(fileProd)
      except OSError:
        continue

    for directory in ['./job', 'job']:
      try:
        shutil.rmtree(directory)
      except Exception:
        continue


# FailoverRequest.py
@patch("LHCbDIRAC.Workflow.Modules.ModuleBase.RequestValidator", side_effect=MagicMock())
class FailoverRequestSuccess(ModulesTestCase):

  #################################################

  def test_execute(self, _patch):

    fr = FailoverRequest(bkClient=bkc_mock, dm=dm_mock)
    fr.jobType = 'merge'
    fr.stepInputData = ['foo', 'bar']
    fr.requestValidator = MagicMock()

    # no errors, no input data
    for wf_cs in copy.deepcopy(wf_commons):
      for s_cs in step_commons:
        self.assertTrue(fr.execute(prod_id, prod_job_id, wms_job_id,
                                   workflowStatus, stepStatus,
                                   wf_cs, s_cs,
                                   step_number, step_id)['OK'])


# RemoveInputData.py
@patch("LHCbDIRAC.Workflow.Modules.ModuleBase.RequestValidator", side_effect=MagicMock())
class RemoveInputDataSuccess(ModulesTestCase):

  #################################################

  def test_execute(self, _patch):

    rid = RemoveInputData(bkClient=bkc_mock, dm=dm_mock)
    # no errors, no input data
    for wf_cs in copy.deepcopy(wf_commons):
      if 'InputData' in wf_cs:
        continue
      for s_cs in step_commons:
        self.assertTrue(rid.execute(prod_id, prod_job_id, wms_job_id,
                                    workflowStatus, stepStatus,
                                    wf_cs, s_cs,
                                    step_number, step_id)['OK'])

    # no errors, input data
    for wf_cs in copy.deepcopy(wf_commons):
      if 'InputData' not in wf_cs:
        continue
      for s_cs in step_commons:
        self.assertTrue(rid.execute(prod_id, prod_job_id, wms_job_id,
                                    workflowStatus, stepStatus,
                                    wf_cs, s_cs,
                                    step_number, step_id)['OK'])


# StepAccounting.py
@patch("LHCbDIRAC.Workflow.Modules.ModuleBase.RequestValidator", side_effect=MagicMock())
class StepAccountingSuccess(ModulesTestCase):

  #################################################

  def test_execute(self, _patch):

    sa = StepAccounting(bkClient=bkc_mock, dm=dm_mock)
    sa.jobType = 'merge'
    sa.stepInputData = ['foo', 'bar']
    sa.siteName = 'DIRAC.Test.ch'

    for wf_cs in copy.deepcopy(wf_commons):
      for s_cs in step_commons:
        self.assertTrue(sa.execute(prod_id, prod_job_id, wms_job_id,
                                   workflowStatus, stepStatus,
                                   wf_cs, s_cs,
                                   step_number, step_id,
                                   self.jobStep_mock, self.xf_o_mock)['OK'])


# UploadLogFile.py
@patch("LHCbDIRAC.Workflow.Modules.ModuleBase.RequestValidator", side_effect=MagicMock())
class UploadLogFileSuccess(ModulesTestCase):

  @patch("LHCbDIRAC.Workflow.Modules.UploadLogFile.getDestinationSEList", side_effect=getDestinationSEListMockCNAF)
  def test_execute(self, _patch, _patch1):

    # no errors, no input data
    #    for wf_commons in copy.deepcopy( wf_commons ):
    #      for step_commons in step_commons:
    #        self.assertTrue( ulf.execute( prod_id, prod_job_id, wms_job_id,
    #                                           workflowStatus, stepStatus,
    #                                           wf_commons, step_commons,
    #                                           step_number, step_id,
    #                                           dm_mock, self.ft_mock,
    #                                           bkc_mock )['OK'] )

    # putStorageDirectory returns False

    rm_mock = copy.deepcopy(dm_mock)
    rm_mock.putStorageDirectory.return_value = {'OK': False, 'Message': 'bih'}
    ft_mock = copy.deepcopy(self.ft_mock)
    ulf = UploadLogFile(bkClient=bkc_mock, dm=dm_mock)
    ulf.request = Request()
    ulf.failoverTransfer = ft_mock
    for wf_cs in copy.deepcopy(wf_commons):
      for s_cs in step_commons:
        self.assertTrue(ulf.execute(prod_id, prod_job_id, 0,
                                    workflowStatus, stepStatus,
                                    wf_cs, s_cs,
                                    step_number, step_id)['OK'])
      # self.assertTrue( ulf.finalize( rm_mock, self.ft_mock )['OK'] )

  @patch("LHCbDIRAC.Workflow.Modules.UploadLogFile.getDestinationSEList", side_effect=getDestinationSEListMockCNAF)
  def test__uploadLogToFailoverSE(self, _patch, _patch1):
    open('foo.txt', 'w').close()
    tarFileName = 'foo.txt'
    ulf = UploadLogFile(bkClient=bkc_mock, dm=dm_mock)
    ulf.request = Request()
    ulf.failoverTransfer = self.ft_mock
    ulf.logLFNPath = '/an/lfn/foo.txt'
    ulf.failoverSEs = ['SE1', 'SE2']
    ulf._uploadLogToFailoverSE(tarFileName)

  @patch("LHCbDIRAC.Workflow.Modules.UploadLogFile.getDestinationSEList", side_effect=getDestinationSEListMockCNAF)
  def test__determinRelevantFiles(self, _patch, _patch1):
    for fileN in ['foo.txt', 'aLog.log', 'aLongLog.log', 'aLongLog.log.gz']:
      try:
        os.remove(fileN)
      except OSError:
        continue

    open('foo.txt', 'w').close()
    open('bar.py', 'w').close()
    open('aLog.log', 'w').close()
    ulf = UploadLogFile(bkClient=bkc_mock, dm=dm_mock)
    res = ulf._determineRelevantFiles()
    self.assertTrue(res['OK'])
    expected = ['foo.txt', 'aLog.log']
    if 'pylint.txt' in os.listdir('.'):
      expected.append('pylint.txt')
    if 'nosetests.xml' in os.listdir('.'):
      expected.append('nosetests.xml')
    self.assertTrue(set(res['Value']) >= set(expected))

    fd = open('aLongLog.log', 'w')
    for _x in range(2500):
      fd.writelines(
          "Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do "
          "eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim "
          "ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut "
          "aliquip ex ea commodo consequat. Duis aute irure dolor in "
          "reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla "
          "pariatur. Excepteur sint occaecat cupidatat non proident, sunt in "
          "culpa qui officia deserunt mollit anim id est laborum"
      )
    fd.close()
    res = ulf._determineRelevantFiles()
    self.assertTrue(res['OK'])
    expected = ['foo.txt', 'aLog.log', 'aLongLog.log']
    if 'pylint.txt' in os.listdir('.'):
      expected.append('pylint.txt')
    if 'nosetests.xml' in os.listdir('.'):
      expected.append('nosetests.xml')
    self.assertTrue(set(res['Value']) >= set(expected))

    open('foo.txt', 'w').close()
    fd = open('aLongLog.log', 'w')
    for _x in range(2500):
      fd.writelines(
          "Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do "
          "eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim "
          "ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut "
          "aliquip ex ea commodo consequat. Duis aute irure dolor in "
          "reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla "
          "pariatur. Excepteur sint occaecat cupidatat non proident, sunt in "
          "culpa qui officia deserunt mollit anim id est laborum"
      )
    fd.close()
    open('bar.py', 'w').close()
    res = ulf._determineRelevantFiles()
    expected = ['foo.txt', 'aLog.log', 'aLongLog.log']
    if 'pylint.txt' in os.listdir('.'):
      expected.append('pylint.txt')
    if 'nosetests.xml' in os.listdir('.'):
      expected.append('nosetests.xml')
    self.assertTrue(res['OK'])
    self.assertTrue(set(res['Value']) >= set(expected))


# UserJobFinalization.py
@patch("LHCbDIRAC.Workflow.Modules.ModuleBase.RequestValidator", side_effect=MagicMock())
class UserJobFinalizationSuccess(ModulesTestCase):

  def test_execute(self, _patch):

    ujf = UserJobFinalization(bkClient=bkc_mock, dm=dm_mock)
    ujf.bkClient = bkc_mock
    ujf.failoverTransfer = self.ft_mock
    ujf.requestValidator = MagicMock()
    ujf.requestValidator.validate.return_value = {'OK': True}

    # no errors, no input data
    for wf_cs in copy.deepcopy(wf_commons):
      wf_cs['TotalSteps'] = step_number
      for s_cs in step_commons:
        self.assertTrue(ujf.execute(prod_id, prod_job_id, wms_job_id,
                                    workflowStatus, stepStatus,
                                    wf_cs, s_cs,
                                    step_number, step_id)['OK'])

    for wf_cs in copy.deepcopy(wf_commons):
      wf_cs['TotalSteps'] = step_number
      for s_cs in step_commons:
        wf_cs['UserOutputData'] = ['i1', 'i2']
        wf_cs['UserOutputSE'] = ['MySE']
        wf_cs['OwnerName'] = 'fstagni'
        open('i1', 'w').close()
        open('i2', 'w').close()
        self.assertTrue(ujf.execute(prod_id, prod_job_id, wms_job_id,
                                    workflowStatus, stepStatus,
                                    wf_cs, s_cs,
                                    step_number, step_id, orderedSEs=['MySE1', 'MySE2'])['OK'])
      os.remove('i1')
      os.remove('i2')

  @patch("LHCbDIRAC.Workflow.Modules.UserJobFinalization.getDestinationSEList", side_effect=getDestinationSEListMock)
  def test__getOrderedSEsList(self, _patch, _patched):

    ujf = UserJobFinalization(bkClient=bkc_mock, dm=dm_mock)

    ujf.userOutputSE = ['userSE']
    res = ujf._getOrderedSEsList()
    self.assertEqual(res, ['userSE'])

    ujf.defaultOutputSE = ['CERN']
    res = ujf._getOrderedSEsList()
    self.assertEqual(res, ['userSE', 'CERN'])

  @patch("LHCbDIRAC.Workflow.Modules.UserJobFinalization.getDestinationSEList",
         side_effect=getDestinationSEListMockCNAF)
  def test__getOrderedSEsListCNAF(self, _patch, _patched):

    ujf = UserJobFinalization(bkClient=bkc_mock, dm=dm_mock)
    res = ujf._getOrderedSEsList()
    self.assertEqual(res, ['CNAF'])


# FileUsage.py
@patch("LHCbDIRAC.Workflow.Modules.ModuleBase.RequestValidator", side_effect=MagicMock())
class FileUsageSuccess(ModulesTestCase):

  #################################################

  def test_execute(self, _patch):

    # no errors, no input files to report
    wfStatus = copy.deepcopy(workflowStatus)
    wfC = copy.deepcopy(wf_commons)
    for wf_cs in wfC:  # copy.deepcopy( wf_commons ):
      for s_cs in step_commons:
        fu = FileUsage(bkClient=bkc_mock, dm=dm_mock)
        self.assertTrue(fu.execute(prod_id, prod_job_id, wms_job_id,
                                   wfStatus, stepStatus,
                                   wf_cs, s_cs,
                                   step_number, step_id)['OK'])
    wfC = copy.deepcopy(wf_commons)
    for wf_cs in wfC:  # copy.deepcopy( wf_commons ):
      wf_cs['ParametricInputData'] = ['LFN:/lhcb/data/2010/EW.DST/00008380/0000/00008380_00000287_1.ew.dst',
                                      'LFN:/lhcb/data/2010/EW.DST/00008380/0000/00008380_00000285_1.ew.dst',
                                      'LFN:/lhcb/data/2010/EW.DST/00008380/0000/00008380_00000281_1.ew.dst']
      for s_cs in step_commons:
        fu = FileUsage(bkClient=bkc_mock, dm=dm_mock)
        self.assertTrue(fu.execute(prod_id, prod_job_id, wms_job_id,
                                   workflowStatus, stepStatus,
                                   wf_cs, s_cs,
                                   step_number, step_id)['OK'])
    wfC = copy.deepcopy(wf_commons)
    for wf_cs in wfC:  # copy.deepcopy( wf_commons ):
      wf_cs['ParametricInputData'] = ['LFN:/lhcb/data/2010/EW.DST/00008380/0000/00008380_00000287_1.ew.dst',
                                      'LFN:/lhcb/data/2010/EW.DST/00008380/0000/00008380_00000285_1.ew.dst',
                                      'LFN:/lhcb/data/2010/PIPPO/00008380/0000/00008380_00000281_1.pippo.dst']
      for s_cs in step_commons:
        fu = FileUsage(bkClient=bkc_mock, dm=dm_mock)
        self.assertTrue(fu.execute(prod_id, prod_job_id, wms_job_id,
                                   workflowStatus, stepStatus,
                                   wf_cs, s_cs,
                                   step_number, step_id)['OK'])

    # workflow status not ok
    wfC = copy.deepcopy(wf_commons)
    for wf_cs in wfC:  # copy.deepcopy( wf_commons ):
      wfStatus = {'OK': False, 'Message': 'Mess'}
      for s_cs in step_commons:
        fu = FileUsage(bkClient=bkc_mock, dm=dm_mock)
        self.assertTrue(fu.execute(prod_id, prod_job_id, wms_job_id,
                                   wfStatus, stepStatus,
                                   wf_cs, s_cs,
                                   step_number, step_id)['OK'])


# FileUsage.py
@patch("LHCbDIRAC.Workflow.Modules.ModuleBase.RequestValidator", side_effect=MagicMock())
class CreateDataFileSuccess(ModulesTestCase):

  #################################################

  def test_execute(self, _patch):

    cdf = CreateDataFile(bkClient=bkc_mock, dm=dm_mock)
    cdf.jobType = 'merge'
    cdf.stepInputData = ['foo', 'bar']

    for wf_cs in copy.deepcopy(wf_commons):
      for s_cs in step_commons:
        self.assertTrue(cdf.execute(prod_id, prod_job_id, wms_job_id,
                                    workflowStatus, stepStatus,
                                    wf_cs, s_cs,
                                    step_number, step_id)['OK'])


if __name__ == '__main__':
  unittest.main()
