# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

import os

import pytest
from pydantic import BaseModel, Extra

from czc.api.v2.constants import CLOUDZERO_TEST_KEY
from czt.unittest import Context


@pytest.fixture(scope='function')
def context(mocker, environ):
    from czc.state_handler import state_handler

    with Context(mocker, state_handler, modules_to_mock=[], extra={'state_handler': state_handler}) as context:
        yield context


class TestEvent(BaseModel, extra=Extra.allow):
    field1: str
    field2: int


@pytest.mark.unit
@pytest.mark.parametrize('with_test_key', [True, False])
def test_state_handler(context, with_test_key):
    test_key_value = 'test-key-value'

    @context.state_handler(TestEvent)
    def handler(event, _, parsed_event: TestEvent):
        if with_test_key:
            assert os.environ[CLOUDZERO_TEST_KEY] == test_key_value
        else:
            assert CLOUDZERO_TEST_KEY not in os.environ

        return parsed_event

    event = {'field1': 'value1', 'field2': 2, **({'__test_key__': test_key_value} if with_test_key else {})}

    result = handler(event, None)

    assert CLOUDZERO_TEST_KEY not in os.environ
    assert result.dict() == event
