# Copyright (C) 2017-2018  Lars Wirzenius
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


import os
import shutil
import tempfile
import unittest


import ick2


class WorkAPITests(unittest.TestCase):

    def setUp(self):
        self.tempdir = tempfile.mkdtemp()
        self.statedir = os.path.join(self.tempdir, 'state/dir')
        self.state = ick2.ControllerState()
        self.state.set_state_directory(self.statedir)

    def tearDown(self):
        shutil.rmtree(self.tempdir)

    def create_project_api(self):
        pipeline = {
            'pipeline': 'build',
            'actions': [
                {'shell': 'step-1', 'where': 'host'},
                {'shell': 'step-2', 'where': 'host'},
            ],
        }

        pipeapi = ick2.PipelineAPI(self.state)
        pipeapi.create(pipeline)

        project = {
            'project': 'foo',
            'parameters': {
                'foo': 'bar',
            },
            'pipelines': ['build'],
        }
        api = ick2.ProjectAPI(self.state)
        api.create(project)
        return api

    def create_worker_api(self):
        worker = {
            'this': 'that',
        }
        claims = {
            'aud': 'asterix',
        }
        api = ick2.WorkerAPI(self.state)
        api.create(worker, claims=claims)
        return api

    def create_work_api(self):
        return ick2.WorkAPI(self.state)

    def test_worker_gets_no_work_when_no_pipeline_is_triggered(self):
        self.create_project_api()
        self.create_worker_api()
        work = self.create_work_api()
        claims = {'aud': 'asterix'}
        self.assertEqual(work.get_work(claims=claims), {})

    def test_worker_gets_work_when_a_pipeline_is_triggered(self):
        projects = self.create_project_api()
        projects.set_status('foo', 'triggered')
        self.create_worker_api()
        work = self.create_work_api()
        expected = {
            'build_id': 'foo/1',
            'build_number': 1,
            'worker': 'asterix',
            'project': 'foo',
            'parameters': {
                'foo': 'bar',
            },
            'step': {
                'action': 'create_workspace',
                'where': 'host',
            },
            'log': '/logs/foo/1',
        }
        claims = {'aud': 'asterix'}
        self.assertEqual(work.get_work(claims=claims), expected)

        # Check we get the same thing twice.
        claims = {'aud': 'asterix'}
        self.assertEqual(work.get_work(claims=claims), expected)

    def test_worker_manager_posts_work_updates(self):
        projects = self.create_project_api()
        projects.set_status('foo', 'triggered')
        self.create_worker_api()
        work = self.create_work_api()

        # Ask for some work.
        expected = {
            'build_id': 'foo/1',
            'build_number': 1,
            'worker': 'asterix',
            'project': 'foo',
            'parameters': {
                'foo': 'bar',
            },
            'step': {
                'action': 'create_workspace',
                'where': 'host',
            },
            'log': '/logs/foo/1',
        }
        claims = {'aud': 'asterix'}
        self.assertEqual(work.get_work(claims=claims), expected)

        # Post a partial update.
        done = {
            'build_id': 'foo/1',
            'worker': 'asterix',
            'project': 'foo',
            'exit_code': None,
            'stdout': 'out',
            'stderr': 'err',
            'timestamp': '2000-01-01T00:00:00',
        }
        work.update_work(done)

        # Ask for work again. We didn't finish the previous step, so
        # should get same thing.
        claims = {'aud': 'asterix'}
        self.assertEqual(work.get_work(claims=claims), expected)

        # Finish the step.
        done['exit_code'] = 0
        work.update_work(done)

        # We should get the next step now.
        expected['step'] = {'shell': 'step-1', 'where': 'host'}
        self.assertEqual(work.get_work(claims=claims), expected)

        # Finish the step.
        done['exit_code'] = 0
        work.update_work(done)

        # We should get the next step now.
        expected['step'] = {'shell': 'step-2', 'where': 'host'}
        self.assertEqual(work.get_work(claims=claims), expected)

        # Finish the step.
        done['exit_code'] = 0
        work.update_work(done)

        # We now get nothing further to do.
        self.assertEqual(work.get_work(claims=claims), {})

        # An pipeline status has changed.
        self.assertEqual(
            projects.get_status('foo'),
            {'status': 'idle'})

    def test_worker_manager_posts_failure(self):
        projects = self.create_project_api()
        projects.set_status('foo', 'triggered')
        self.create_worker_api()
        work = self.create_work_api()

        # Ask for some work.
        expected = {
            'build_id': 'foo/1',
            'build_number': 1,
            'worker': 'asterix',
            'project': 'foo',
            'parameters': {
                'foo': 'bar',
            },
            'step': {
                'action': 'create_workspace',
                'where': 'host',
            },
            'log': '/logs/foo/1',
        }
        claims = {'aud': 'asterix'}
        self.assertEqual(work.get_work(claims=claims), expected)

        # Post a partial update.
        done = {
            'build_id': 'foo/1',
            'worker': 'asterix',
            'project': 'foo',
            'exit_code': 1,
            'stdout': 'out',
            'stderr': 'err',
            'timestamp': '2000-01-01T00:00:00',
        }
        work.update_work(done)

        # Ask for work again.
        claims = {'aud': 'asterix'}
        self.assertEqual(work.get_work(claims=claims), {})

        # And project status has changed.
        self.assertEqual(
            projects.get_status('foo'),
            {'status': 'idle'})