# Copyright (C) 2017-2019  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 copy
import os
import unittest


import ick2


class WorkAPITests(unittest.TestCase):

    def setUp(self):
        self.state = ick2.MemoryStore()
        self.claims = None

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

        getter = DummyTokenGetter()
        pipeapi = ick2.PipelineAPI(self.state)
        pipeapi.set_token_getter(getter)
        pipeapi.create(pipeline)

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

    def create_worker_api(self):
        worker = {
            'doing': {},
        }
        self.claims = {
            'aud': 'asterix',
        }
        getter = DummyTokenGetter()
        api = ick2.WorkerAPI(self.state)
        api.set_token_getter(getter)
        api.create(worker, claims=self.claims)
        return api

    def create_work_api(self):
        getter = DummyTokenGetter()
        api = ick2.WorkAPI(self.state)
        api.set_token_getter(getter)
        return api

    def test_worker_gets_no_work_when_no_builds_have_been_triggered(self):
        self.create_project_api()
        self.create_worker_api()
        work = self.create_work_api()
        self.assertEqual(work.get_work(claims=self.claims), {})

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

        # Check we get the same thing twice.
        self.assertEqual(work.get_work(claims=self.claims), expected)

    def test_worker_manager_posts_work_updates(self):
        # Define the actions we expect to get from the controller.
        # They're mostly identical so we copy and change what needs to
        # be changed.

        action_1 = {
            'build_id': 'foo/1',
            'build_number': 1,
            'worker': 'asterix',
            'project': 'foo',
            'parameters': {
                'foo': 'bar',
            },
            'action_id': '1',
            'step': {
                'action': 'create_workspace',
                'where': 'host',
            },
            'log': '/logs/foo/1',
        }

        action_2 = copy.deepcopy(action_1)
        action_2.update({
            'action_id': '2',
            'step': {
                'shell': 'step-1',
                'where': 'host',
            },
        })

        action_3 = copy.deepcopy(action_1)
        action_3.update({
            'action_id': '3',
            'step': {
                'shell': 'step-2',
                'where': 'host',
            },
        })

        action_4 = copy.deepcopy(action_1)
        action_4.update({
            'action_id': '4',
            'step': {
                'action': 'notify',
            },
        })

        # Define the work updates we will be sending to indicate
        # progress on executing the actions above. Again, they're
        # mostly identical.

        done_1_partial = {
            'build_id': 'foo/1',
            'action_id': '1',
            'worker': 'asterix',
            'project': 'foo',
            'exit_code': None,
            'stdout': 'out',
            'stderr': 'err',
            'timestamp': '2000-01-01T00:00:00',
        }
        done_1 = copy.deepcopy(done_1_partial)
        done_1.update({
            'exit_code': 0,
        })
        done_2 = copy.deepcopy(done_1)
        done_2.update({
            'action_id': '2',
        })
        done_3 = copy.deepcopy(done_1)
        done_3.update({
            'action_id': '3',
        })
        done_4 = copy.deepcopy(done_1)
        done_4.update({
            'action_id': '4',
        })

        # Set up the various API objects.
        projects = self.create_project_api()
        self.create_worker_api()
        work = self.create_work_api()

        # No builds have been triggered, nothing to do.
        self.assertEqual(work.get_work(claims=self.claims), {})

        # Trigger a build.
        projects.trigger_project('foo')

        # Get the first action.
        self.assertEqual(work.get_work(claims=self.claims), action_1)

        # Post a partial update.
        work.update_work(done_1_partial)

        # Ask for work again. We didn't finish the previous step, so
        # should get same thing.
        self.assertEqual(work.get_work(claims=self.claims), action_1)

        # Finish the step.
        work.update_work(done_1)

        # We should get the next step now.
        self.assertEqual(work.get_work(claims=self.claims), action_2)

        # Finish the step.
        work.update_work(done_2)

        # We should get the final actual step now.
        self.assertEqual(work.get_work(claims=self.claims), action_3)

        # Finish the step.
        work.update_work(done_3)

        # We should get the notification step now.
        self.assertEqual(work.get_work(claims=self.claims), action_4)

        # Finish the step.
        work.update_work(done_4)

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

    def test_worker_manager_posts_failure(self):
        projects = self.create_project_api()
        projects.trigger_project('foo')
        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',
            },
            'action_id': '1',
            'step': {
                'action': 'create_workspace',
                'where': 'host',
            },
            'log': '/logs/foo/1',
        }
        self.assertEqual(work.get_work(claims=self.claims), expected)

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

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

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

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


class DummyTokenGetter:

    def get_token(self):
        return 'DUMMY.TOKEN'