# -*- coding: utf-8 -*-
# Copyright 2011 Matthew Tai
# Copyright 2012 Yelp
# Copyright 2013 Yelp and Lyft
# Copyright 2014 Marc Abramowitz
# Copyright 2015-2016 Yelp
#
# 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.
"""Tests for InlineMRJobRunner"""
import gzip
import os
import os.path
from io import BytesIO

from mrjob import conf
from mrjob.fs.base import Filesystem
from mrjob.inline import InlineMRJobRunner
from mrjob.job import MRJob
from mrjob.protocol import JSONValueProtocol
from mrjob.sim import _error_on_bad_paths
from mrjob.step import MRStep
from tests.mr_no_mapper import MRNoMapper
from tests.mr_test_cmdenv import MRTestCmdenv
from tests.mr_test_jobconf import MRTestJobConf
from tests.mr_test_per_step_jobconf import MRTestPerStepJobConf
from tests.mr_two_step_job import MRTwoStepJob
from tests.mr_word_count import MRWordCount
from tests.py2 import TestCase
from tests.py2 import mock
from tests.py2 import patch
from tests.sandbox import EmptyMrjobConfTestCase
from tests.sandbox import SandboxedTestCase


class InlineMRJobRunnerEndToEndTestCase(SandboxedTestCase):

    def test_end_to_end(self):
        # read from STDIN, a regular file, and a .gz
        stdin = BytesIO(b'foo\nbar\n')

        input_path = os.path.join(self.tmp_dir, 'input')
        with open(input_path, 'w') as input_file:
            input_file.write('bar\nqux\n')

        input_gz_path = os.path.join(self.tmp_dir, 'input.gz')
        input_gz = gzip.GzipFile(input_gz_path, 'wb')
        input_gz.write(b'foo\n')
        input_gz.close()

        mr_job = MRTwoStepJob(
            ['--runner', 'inline', '-', input_path, input_gz_path])
        mr_job.sandbox(stdin=stdin)

        local_tmp_dir = None
        results = []

        with mr_job.make_runner() as runner:
            assert isinstance(runner, InlineMRJobRunner)
            runner.run()

            for line in runner.stream_output():
                key, value = mr_job.parse_output_line(line)
                results.append((key, value))

            local_tmp_dir = runner._get_local_tmp_dir()
            assert os.path.exists(local_tmp_dir)

        # make sure cleanup happens
        assert not os.path.exists(local_tmp_dir)

        self.assertEqual(sorted(results),
                         [(1, 'qux'), (2, 'bar'), (2, 'foo'), (5, None)])

    def test_missing_input(self):
        runner = InlineMRJobRunner(input_paths=['/some/bogus/file/path'])
        self.assertRaises(Exception, runner._run)


class InlineMRJobRunnerCmdenvTest(EmptyMrjobConfTestCase):

    def test_cmdenv(self):
        import logging
        logging.basicConfig()
        # make sure previous environment is preserved
        os.environ['SOMETHING'] = 'foofoofoo'
        old_env = os.environ.copy()

        mr_job = MRTestCmdenv(['--runner', 'inline', '--cmdenv=FOO=bar'])
        mr_job.sandbox(stdin=BytesIO(b'foo\n'))

        results = []

        with mr_job.make_runner() as runner:
            assert isinstance(runner, InlineMRJobRunner)
            runner.run()

            for line in runner.stream_output():
                key, value = mr_job.parse_output_line(line)
                results.append((key, value))

        self.assertEqual(sorted(results),
                         [('FOO', 'bar'), ('SOMETHING', 'foofoofoo')])

        # make sure we revert back
        self.assertEqual(old_env, os.environ)


# this doesn't need to be in its own file because it'll be run inline
class MRIncrementerJob(MRJob):
    """A terribly silly way to add a positive integer to values."""

    INPUT_PROTOCOL = JSONValueProtocol
    OUTPUT_PROTOCOL = JSONValueProtocol

    def configure_options(self):
        super(MRIncrementerJob, self).configure_options()

        self.add_passthrough_option('--times', type='int', default=1)

    def mapper(self, _, value):
        yield None, value + 1

    def steps(self):
        return [MRStep(mapper=self.mapper)] * self.options.times


class MRCustomFileOptionJob(MRJob):
    """ A simple MRJob that uses the input file option."""

    INPUT_PROTOCOL = JSONValueProtocol
    OUTPUT_PROTOCOL = JSONValueProtocol

    multiplier = 1

    def configure_options(self):
        super(MRCustomFileOptionJob, self).configure_options()
        self.add_file_option('--platform_file')

    def mapper_init(self):
        with open(self.options.platform_file) as f:
            self.multiplier = int(f.read())

    def mapper(self, _, value):
        yield None, value * self.multiplier


class InlineRunnerStepsTestCase(EmptyMrjobConfTestCase):
    # make sure file options get passed to --steps in inline mode

    def test_adding_2(self):
        mr_job = MRIncrementerJob(['-r', 'inline', '--times', '2'])
        mr_job.sandbox(stdin=BytesIO(b'0\n1\n2\n'))

        self.assertEqual(len(mr_job.steps()), 2)

        with mr_job.make_runner() as runner:
            assert isinstance(runner, InlineMRJobRunner)
            self.assertEqual(runner._get_steps(), [
                {
                    'type': 'streaming',
                    'mapper': {
                        'type': 'script',
                    }
                },
                {
                    'type': 'streaming',
                    'mapper': {
                        'type': 'script',
                    }
                },
            ])

            runner.run()

            output = sorted(mr_job.parse_output_line(line)[1]
                            for line in runner.stream_output())

            self.assertEqual(output, [2, 3, 4])


class MRJobFileOptionsTestCase(SandboxedTestCase):

    def setUp(self):
        super(MRJobFileOptionsTestCase, self).setUp()

        self.input_file_path = os.path.join(self.tmp_dir, 'input_file.txt')
        with open(self.input_file_path, 'wb') as f:
            f.write(b'2\n')

    def test_with_input_file_option(self):
        mr_job = MRCustomFileOptionJob(
            ['-r', 'inline', '--platform_file', self.input_file_path])
        mr_job.sandbox(stdin=BytesIO(b'1\n'))

        with mr_job.make_runner() as runner:
            runner.run()
            output = sorted(mr_job.parse_output_line(line)[1]
                            for line in runner.stream_output())

            self.assertEqual(output, [2])


class NoMRJobConfTestCase(TestCase):

    def test_no_mrjob_confs(self):
        with patch.object(
                conf, '_expanded_mrjob_conf_path', return_value=None):

            mr_job = MRIncrementerJob(['-r', 'inline', '--times', '2'])
            mr_job.sandbox(stdin=BytesIO(b'0\n1\n2\n'))

            with mr_job.make_runner() as runner:
                runner.run()
                output = sorted(mr_job.parse_output_line(line)[1]
                                for line in runner.stream_output())
                self.assertEqual(output, [2, 3, 4])


class InlineMRJobRunnerJobConfTestCase(SandboxedTestCase):

    # this class is also used to test local mode
    RUNNER = 'inline'

    def test_input_files_and_setting_number_of_tasks(self):
        input_path = os.path.join(self.tmp_dir, 'input')
        with open(input_path, 'wb') as input_file:
            input_file.write(b'bar\nqux\nfoo\n')

        input_gz_path = os.path.join(self.tmp_dir, 'input.gz')
        input_gz = gzip.GzipFile(input_gz_path, 'wb')
        input_gz.write(b'foo\n')
        input_gz.close()

        mr_job = MRWordCount(['-r', self.RUNNER,
                              '--jobconf=mapred.map.tasks=3',
                              '--jobconf=mapred.reduce.tasks=3',
                              input_path, input_gz_path])
        mr_job.sandbox()

        results = []

        with mr_job.make_runner() as runner:
            runner.run()

            for line in runner.stream_output():
                key, value = mr_job.parse_output_line(line)
                results.append((key, value))

            self.assertEqual(runner.counters()[0]['count']['combiners'], 3)

        self.assertEqual(sorted(results),
                         [(input_path, 3), (input_gz_path, 1)])

    def test_jobconf_simulated_by_runner(self):
        input_path = os.path.join(self.tmp_dir, 'input')
        with open(input_path, 'wb') as input_file:
            input_file.write(b'foo\n')

        upload_path = os.path.join(self.tmp_dir, 'upload')
        with open(upload_path, 'wb') as upload_file:
            upload_file.write(b'PAYLOAD')

        # use --no-bootstrap-mrjob so we don't have to worry about
        # mrjob.tar.gz and the setup wrapper script
        self.add_mrjob_to_pythonpath()
        mr_job = MRTestJobConf(['-r', self.RUNNER,
                                '--no-bootstrap-mrjob',
                                '--jobconf=user.defined=something',
                                '--jobconf=mapred.map.tasks=1',
                                '--file', upload_path,
                               input_path])

        mr_job.sandbox()

        results = {}

        # between the single line of input and setting mapred.map.tasks to 1,
        # we should be restricted to only one task, which will give more
        # predictable results

        with mr_job.make_runner() as runner:
            script_path = runner._script_path

            runner.run()

            for line in runner.stream_output():
                key, value = mr_job.parse_output_line(line)
                results[key] = value

        working_dir = results['mapreduce.job.local.dir']
        self.assertEqual(working_dir,
                         os.path.join(runner._get_local_tmp_dir(),
                                      'job_local_dir', '0', 'mapper', '0'))

        self.assertEqual(results['mapreduce.job.cache.archives'], '')
        expected_cache_files = (
            script_path + '#mr_test_jobconf.py',
            upload_path + '#upload')
        self.assertEqual(
            sorted(results['mapreduce.job.cache.files'].split(',')),
            sorted(expected_cache_files))
        self.assertEqual(results['mapreduce.job.cache.local.archives'], '')
        expected_local_files = (
            os.path.join(working_dir, 'mr_test_jobconf.py'),
            os.path.join(working_dir, 'upload'))
        self.assertEqual(
            sorted(results['mapreduce.job.cache.local.files'].split(',')),
            sorted(expected_local_files))
        self.assertEqual(results['mapreduce.job.id'], runner._job_key)

        self.assertEqual(results['mapreduce.map.input.file'], input_path)
        self.assertEqual(results['mapreduce.map.input.length'], '4')
        self.assertEqual(results['mapreduce.map.input.start'], '0')
        self.assertEqual(results['mapreduce.task.attempt.id'],
                         'attempt_%s_mapper_00000_0' % runner._job_key)
        self.assertEqual(results['mapreduce.task.id'],
                         'task_%s_mapper_00000' % runner._job_key)
        self.assertEqual(results['mapreduce.task.ismap'], 'true')
        self.assertEqual(results['mapreduce.task.output.dir'],
                         runner._output_dir)
        self.assertEqual(results['mapreduce.task.partition'], '0')
        self.assertEqual(results['user.defined'], 'something')

    def test_per_step_jobconf(self):
        mr_job = MRTestPerStepJobConf([
            '-r', self.RUNNER, '--jobconf', 'user.defined=something'])
        mr_job.sandbox()

        results = {}

        with mr_job.make_runner() as runner:
            runner.run()

            for line in runner.stream_output():
                key, value = mr_job.parse_output_line(line)
                results[tuple(key)] = value

        # user.defined gets re-defined in the second step
        self.assertEqual(results[(0, 'user.defined')], 'something')
        self.assertEqual(results[(1, 'user.defined')], 'nothing')

    def test_per_step_jobconf_can_set_number_of_tasks(self):
        mr_job = MRTestPerStepJobConf([
            '-r', self.RUNNER, '--jobconf', 'mapred.map.tasks=2',
        ])
        # need at least two items of input to get two map tasks
        mr_job.sandbox(BytesIO(b'foo\nbar\n'))

        with mr_job.make_runner() as runner:
            runner.run()

            # sanity test: --jobconf should definitely work
            self.assertEqual(runner.counters()[0]['count']['mapper_init'], 2)
            # the job sets its own mapred.map.tasks to 4 for the 2nd step
            self.assertEqual(runner.counters()[1]['count']['mapper_init'], 4)


class InlineMRJobRunnerNoMapperTestCase(SandboxedTestCase):

    RUNNER = 'inline'

    # tests #1141. Also used by local mapper

    def test_step_with_no_mapper(self):
        mr_job = MRNoMapper(['-r', self.RUNNER])

        mr_job.sandbox(stdin=BytesIO(
            b'one fish two fish\nred fish blue fish\n'))

        with mr_job.make_runner() as runner:
            runner.run()

            results = [mr_job.parse_output_line(line)
                       for line in runner.stream_output()]

            self.assertEqual(sorted(results),
                             [(1, ['blue', 'one', 'red', 'two']),
                              (4, ['fish'])])


class InlineMRJobRunnerFSTestCase(SandboxedTestCase):

    RUNNER_CLASS = InlineMRJobRunner

    def setUp(self):
        super(InlineMRJobRunnerFSTestCase, self).setUp()
        self.runner = self.RUNNER_CLASS()

    def test_can_handle_paths(self):
        self.assertEqual(
            self.runner.fs.exists(os.path.join(self.tmp_dir, 'foo')), False)

    def test_cant_handle_uris(self):
        self.assertRaises(IOError, self.runner.fs.ls, 's3://walrus/foo')


class ErrorOnBadPathsTestCase(TestCase):

    def setUp(self):
        self.fs = mock.create_autospec(Filesystem)
        self.paths = ['/one', '/two' '/three/*']

    def test_with_paths(self):
        _error_on_bad_paths(self.fs, self.paths)
        self.fs.exists.assert_called_once_with(self.paths[0])

    def test_no_paths(self):
        self.fs.exists.return_value = False
        self.assertRaises(ValueError, _error_on_bad_paths, self.fs, self.paths)
