# -*- coding: utf-8 -*-
#
# Copyright 2012-2015 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import subprocess
import tempfile

import luigi
from helpers import unittest
from luigi.contrib.pig import PigJobError, PigJobTask
from mock import patch

import pytest


class SimpleTestJob(PigJobTask):
    def output(self):
        return luigi.LocalTarget('simple-output')

    def pig_script_path(self):
        return "my_simple_pig_script.pig"


class ComplexTestJob(PigJobTask):
    def output(self):
        return luigi.LocalTarget('complex-output')

    def pig_script_path(self):
        return "my_complex_pig_script.pig"

    def pig_env_vars(self):
        return {'PIG_CLASSPATH': '/your/path'}

    def pig_properties(self):
        return {'pig.additional.jars': '/path/to/your/jar'}

    def pig_parameters(self):
        return {'YOUR_PARAM_NAME': 'Your param value'}

    def pig_options(self):
        return ['-x', 'local']


@pytest.mark.apache
class SimplePigTest(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

    @patch('subprocess.Popen')
    def test_run__success(self, mock):
        arglist_result = []
        p = subprocess.Popen
        subprocess.Popen = _get_fake_Popen(arglist_result, 0)
        try:
            job = SimpleTestJob()
            job.run()
            self.assertEqual([['/usr/share/pig/bin/pig', '-f', 'my_simple_pig_script.pig']], arglist_result)
        finally:
            subprocess.Popen = p

    @patch('subprocess.Popen')
    def test_run__fail(self, mock):
        arglist_result = []
        p = subprocess.Popen
        subprocess.Popen = _get_fake_Popen(arglist_result, 1)
        try:
            job = SimpleTestJob()
            job.run()
            self.assertEqual([['/usr/share/pig/bin/pig', '-f', 'my_simple_pig_script.pig']], arglist_result)
        except PigJobError as e:
            p = e
            self.assertEqual('stderr', p.err)
        else:
            self.fail("Should have thrown PigJobError")
        finally:
            subprocess.Popen = p


@pytest.mark.apache
class ComplexPigTest(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

    @patch('subprocess.Popen')
    def test_run__success(self, mock):
        arglist_result = []
        p = subprocess.Popen
        subprocess.Popen = _get_fake_Popen(arglist_result, 0)

        with tempfile.NamedTemporaryFile(delete=False) as param_file_mock, \
                tempfile.NamedTemporaryFile(delete=False) as prop_file_mock, \
                patch('luigi.contrib.pig.tempfile.NamedTemporaryFile',
                      side_effect=[param_file_mock, prop_file_mock]):
            try:
                job = ComplexTestJob()
                job.run()
                self.assertEqual([['/usr/share/pig/bin/pig', '-x', 'local',
                                   '-param_file', param_file_mock.name,
                                   '-propertyFile', prop_file_mock.name,
                                   '-f', 'my_complex_pig_script.pig']],
                                 arglist_result)

                # Check param file
                with open(param_file_mock.name) as pparams_file:
                    pparams = pparams_file.readlines()
                    self.assertEqual(1, len(pparams))
                    self.assertEqual('YOUR_PARAM_NAME=Your param value\n', pparams[0])

                # Check property file
                with open(prop_file_mock.name) as pprops_file:
                    pprops = pprops_file.readlines()
                    self.assertEqual(1, len(pprops))
                    self.assertEqual('pig.additional.jars=/path/to/your/jar\n', pprops[0])
            finally:
                subprocess.Popen = p

    @patch('subprocess.Popen')
    def test_run__fail(self, mock):
        arglist_result = []
        p = subprocess.Popen
        subprocess.Popen = _get_fake_Popen(arglist_result, 1)

        with tempfile.NamedTemporaryFile(delete=False) as param_file_mock, \
                tempfile.NamedTemporaryFile(delete=False) as prop_file_mock, \
                patch('luigi.contrib.pig.tempfile.NamedTemporaryFile',
                      side_effect=[param_file_mock, prop_file_mock]):
            try:
                job = ComplexTestJob()
                job.run()
            except PigJobError as e:
                p = e
                self.assertEqual('stderr', p.err)
                self.assertEqual([['/usr/share/pig/bin/pig', '-x', 'local',
                                   '-param_file', param_file_mock.name,
                                   '-propertyFile', prop_file_mock.name, '-f',
                                   'my_complex_pig_script.pig']],
                                 arglist_result)

                # Check param file
                with open(param_file_mock.name) as pparams_file:
                    pparams = pparams_file.readlines()
                    self.assertEqual(1, len(pparams))
                    self.assertEqual('YOUR_PARAM_NAME=Your param value\n', pparams[0])

                # Check property file
                with open(prop_file_mock.name) as pprops_file:
                    pprops = pprops_file.readlines()
                    self.assertEqual(1, len(pprops))
                    self.assertEqual('pig.additional.jars=/path/to/your/jar\n', pprops[0])
            else:
                self.fail("Should have thrown PigJobError")
            finally:
                subprocess.Popen = p


def _get_fake_Popen(arglist_result, return_code, *args, **kwargs):
    def Popen_fake(arglist, shell=None, stdout=None, stderr=None, env=None, close_fds=True):
        arglist_result.append(arglist)

        class P:
            number_of_process_polls = 5

            def __init__(self):
                self._process_polls_left = self.number_of_process_polls

            def wait(self):
                pass

            def poll(self):
                if self._process_polls_left:
                    self._process_polls_left -= 1
                    return None

                return 0

            def communicate(self):
                return 'end'

            def env(self):
                return self.env

        p = P()
        p.returncode = return_code

        p.stderr = tempfile.TemporaryFile()
        p.stdout = tempfile.TemporaryFile()

        p.stdout.write(b'stdout')
        p.stderr.write(b'stderr')

        # Reset temp files so the output can be read.
        p.stdout.seek(0)
        p.stderr.seek(0)

        return p

    return Popen_fake
