import unittest
from unittest.mock import Mock

from db_schema.schemas.exploration import DisassembleContentStatus
from parameterized import parameterized

from exp_sf_runner_lambda.entities import LambdaConfig
from exp_sf_runner_lambda.pool import Pool


class LambdaContextMock:
    aws_request_id = 'foo'
    initial_time = 100 * 1000

    def get_remaining_time_in_millis(self):
        self.initial_time -= 5000
        return self.initial_time


class PoolTestCase(unittest.TestCase):
    def setUp(self) -> None:
        self.pg_repo = Mock()
        self.lambda_cfg = LambdaConfig(
            sf_arn='SF_ARN',
            max_active_runs_dbx=3,
            max_active_runs_total=5,
            dest_bucket='DEST_BUCKET',
            exp_config_path='EXP_CONFIG_PATH',
            max_filesize_not_split=500,
            retry_failed_limit=5,
            trace_id='trace_id',
            optional_config={'RDS_SECRET_KEY': 'test/key'}
        )
        self.sf_service = Mock()
        self.pool = Pool(
            logger=Mock(),
            pg_repository=self.pg_repo,
            config=self.lambda_cfg,
            sf_service=self.sf_service,
            lambda_context=LambdaContextMock()
        )

    @parameterized.expand([
        (1, 5),
        (0, 6),
        (3, 0),
    ])
    def test_pool_run(self, currently_in_progress, start_sf_call_count):
        self.pg_repo.get_active_count = lambda **kwargs: currently_in_progress
        self.pg_repo.get_queued = lambda: [
            DisassembleContentStatus(disassemble_content_status_id=1, ),
            DisassembleContentStatus(disassemble_content_status_id=2, ),
            DisassembleContentStatus(disassemble_content_status_id=3, ),
            DisassembleContentStatus(disassemble_content_status_id=4, ),
            DisassembleContentStatus(disassemble_content_status_id=5, ),
            DisassembleContentStatus(disassemble_content_status_id=6, ),
            DisassembleContentStatus(disassemble_content_status_id=7, ),
            DisassembleContentStatus(disassemble_content_status_id=8, ),
        ]

        self.sf_service.start_sf.side_effect = [
            (True, True),  # sf started, dbx used
            (False, False),  # sf not started(something failed)
            (True, False),  # sf started, dbx not used
            (True, False),  # sf started, dbx not used
            (True, True),  # sf started, dbx used
            (True, True),  # sf started, dbx used
            (True, True),  # sf started, dbx used
        ]
        self.pool.run()
        self.assertEqual(self.sf_service.start_sf.call_count, start_sf_call_count)