# Copyright (c) 2021 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
from decimal import Decimal
from unittest.mock import ANY, Mock, call
from uuid import uuid4

import pytest
from boto3.dynamodb.conditions import Attr, Key
from botocore.exceptions import ClientError

from czc import dynamodb as dyn
from czc import feature_data
from czc.feature_data import MAX_RESOLVE_DATA_TRIES, S3_OBJECT_STORE_ROOT, FeatureDataConnection
from czt.unittest import Context, MockEx

EXCEED_MAX_SIZE_ERROR = ClientError(
    error_response=dict(Error=dict(Code='ValidationException', Message='Item size has exceeded the maximum allowed size')),
    operation_name='PutItem',
)
NOT_FOUND_ERROR = ClientError(
    error_response=dict(Error=dict(Code='NoSuchKey', Message='The specified key does not exist.')),
    operation_name='GetObject',
)


@pytest.fixture(scope='function')
def context(mocker):
    extra = dict(conn=FeatureDataConnection(_table=Mock(), _s3_bucket=str(uuid4())), org_id=str(uuid4()), keyspace=str(uuid4()))
    with Context(mocker, feature_data, ['dyn', 's3', 'uuid', 's3_client'], extra=extra) as context:
        os.environ['NAMESPACE'] = 'ns'
        context.mock_dyn.prefix.side_effect = lambda p, x: f'{p}#{x}'
        context.mock_dyn.unfix.side_effect = lambda p, x: x.removeprefix(f'{p}#')
        context.mock_dyn.compose_key.side_effect = lambda *x: '#'.join(x)
        context.mock_dyn.type_for_dynamo.side_effect = dyn.type_for_dynamo
        context.mock_uuid.uuid4.return_value = 'UUID'
        mocker.patch('retrying.time.sleep')

        yield context


@pytest.fixture(scope='function')
def context_dynamo(mocker, mock_table):
    _table = mock_table

    def load_data(data: list[dict]):
        for item in data:
            _table.put_item(Item=item)

    extra = {
        'conn': FeatureDataConnection(_table=_table, _s3_bucket=str(uuid4())),
        'org_id': str(uuid4()),
        'keyspace': str(uuid4()),
        'load_data': load_data,
    }
    with Context(mocker, feature_data, ['s3'], extra=extra) as context:
        os.environ['NAMESPACE'] = 'ns'
        yield context


@pytest.mark.unit
@pytest.mark.parametrize('table_name,bucket_name', [('standard-table', 'standard-bucket')])
def test_connect_should_use_standard_table_and_bucket_names(context, table_name, bucket_name):
    expected_table = Mock()
    MockEx(context.mock_s3.get_bucket_name).if_called_with().return_value(bucket_name)
    MockEx(context.mock_dyn.get_table_name).if_called_with().return_value(table_name)
    MockEx(context.mock_dyn.table).if_called_with(table_name).return_value(expected_table)

    conn = feature_data.connect()
    assert conn._table == expected_table
    assert conn._s3_bucket == bucket_name


@pytest.mark.unit
@pytest.mark.parametrize(
    'org_id,keyspace,keys,expected_dynamo_key',
    [
        ('org1', 'ks', [], dict(PK='organization#org1', SK='ks')),
        ('org2', 'obj', ['k1', 'k2'], dict(PK='organization#org2', SK='obj#k1#k2')),
    ],
)
def test_get_object_should_use_correct_dynamo_key(context, org_id, keyspace, keys, expected_dynamo_key):
    feature_data.get_object(context.conn, org_id, keyspace, *keys)

    context.conn._table.get_item.assert_called_once_with(Key=expected_dynamo_key, ConsistentRead=True)


@pytest.mark.unit
@pytest.mark.parametrize('dynamo_data,expected_result', [(None, {}), ({}, {}), ({'val': 1}, {'val': 1})])
def test_get_object_should_return_dynamo_data(context, dynamo_data, expected_result):
    MockEx(context.conn._table.get_item).if_called_with_anything().return_value(
        dict(Item=dict(data=dynamo_data)) if dynamo_data is not None else dict()
    )
    result = feature_data.get_object(context.conn, context.org_id, context.keyspace)

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'dynamo_data,expected_path,expected_result',
    [('REF:path/to/data', 'path/to/data', {'data': 1}), ('Invalid string', '', ValueError)],
)
def test_get_object_should_return_s3_data_for_ref(context, dynamo_data, expected_path, expected_result):
    MockEx(context.conn._table.get_item).if_called_with_anything().return_value(
        dict(Item=dict(data=dynamo_data)) if dynamo_data is not None else dict()
    )
    MockEx(context.mock_s3.get_json_object).if_called_with(
        context.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{expected_path}'
    ).return_value(expected_result)

    def get_object():
        return feature_data.get_object(context.conn, context.org_id, context.keyspace)

    if expected_result is ValueError:
        pytest.raises(ValueError, get_object)
    else:
        assert get_object() == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'org_id,keyspace,keys,expected_dynamo_condition',
    [
        ('org1', 'ks', [], Key('PK').eq('organization#org1') & Key('SK').begins_with('ks#')),
        ('org2', 'obj', ['k1', 'k2'], Key('PK').eq('organization#org2') & Key('SK').begins_with('obj#k1#k2#')),
    ],
)
def test_list_objects_should_use_correct_dynamo_condition(context, org_id, keyspace, keys, expected_dynamo_condition):
    list(feature_data.list_objects(context.conn, org_id, keyspace, *keys))

    context.mock_dyn.query.assert_called_once_with(
        context.conn._table, dict(KeyConditionExpression=expected_dynamo_condition, ConsistentRead=True)
    )


@pytest.mark.unit
@pytest.mark.parametrize(
    'dynamo_data,expected_result',
    [
        ([], []),
        (
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': {'val': 1}},
                {'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': {'val': 2}},
            ],
            [{'val': 1}, {'val': 2}],
        ),
    ],
)
def test_list_objects_should_return_dynamo_data(context, dynamo_data, expected_result):
    MockEx(context.mock_dyn.query).if_called_with_anything().return_value(dynamo_data)
    result = list(feature_data.list_objects(context.conn, context.org_id, context.keyspace))

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'dynamo_data,mock_s3_data,expected_result',
    [
        (
            [{'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': 'REF:path/to/data'}],
            {'path/to/data': {'data': 1}},
            [{'data': 1}],
        ),
        (
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': 'REF:path/to/data'},
                {'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': 'REF:path2'},
            ],
            {'path/to/data': {'data': 1}, 'path2': {'data': 2}},
            [{'data': 1}, {'data': 2}],
        ),
        (
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': 'REF:path/to/data'},
                {'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': {'data': 2}},
            ],
            {'path/to/data': {'data': 1}},
            [{'data': 1}, {'data': 2}],
        ),
        ([{'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': 'Invalid string'}], {}, ValueError),
    ],
)
def test_list_objects_should_return_s3_data_for_ref(context, dynamo_data, mock_s3_data, expected_result):
    MockEx(context.mock_dyn.query).if_called_with_anything().return_value(dynamo_data)
    for path, data in mock_s3_data.items():
        MockEx(context.mock_s3.get_json_object).if_called_with(
            context.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}'
        ).return_value(data)

    def list_objects():
        return list(feature_data.list_objects(context.conn, context.org_id, context.keyspace))

    if expected_result is ValueError:
        pytest.raises(ValueError, list_objects)
    else:
        assert list_objects() == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'initial_dynamo_data,initial_s3,updated_dynamo_data,update_s3,expected_list',
    [
        pytest.param(
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': {'value': 'IN DYNAMO'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': 'REF:initial/data'},
            ],
            {'initial/data': 'INITIAL'},
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': 'REF:updated/data'},
            ],
            {'updated/data': 'UPDATED'},
            [{'value': 'UPDATED'}, {'value': 'IN DYNAMO'}],
            id='update_s3_ref_to_other_s3_ref',
        ),
        pytest.param(
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': {'value': 'IN DYNAMO'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': 'REF:initial/data'},
            ],
            {'initial/data': 'INITIAL'},
            [{'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': {'value': 'BACK IN DYNAMO'}}],
            {},
            [{'value': 'BACK IN DYNAMO'}, {'value': 'IN DYNAMO'}],
            id='update_s3_ref_to_no_longer_in_s3',
        ),
    ],
)
def test_list_objects_should_reload_changed_s3_references(
    context_dynamo, initial_dynamo_data, initial_s3, updated_dynamo_data, update_s3, expected_list
):
    def mock_s3(remove: dict, add: dict):
        for path in remove.keys():
            MockEx(context_dynamo.mock_s3.get_json_object).if_called_with(
                context_dynamo.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}'
            ).raises(NOT_FOUND_ERROR)
        for path, value in add.items():
            # new reference exists now
            MockEx(context_dynamo.mock_s3.get_json_object).if_called_with(
                context_dynamo.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}'
            ).return_value({'value': value})

    mock_s3({}, initial_s3)
    context_dynamo.load_data(initial_dynamo_data)

    # get the list generator
    gen = feature_data.list_objects(context_dynamo.conn, 'foo', 'keyspace')

    # change the dynamo reference and s3 contents
    mock_s3(initial_s3, update_s3)
    context_dynamo.load_data(updated_dynamo_data)

    # resolving the generator to a list should get the updated data
    assert list(gen) == expected_list


@pytest.mark.unit
@pytest.mark.parametrize(
    'initial_dynamo_data,s3_data,removed_dynamo_data,expected_list',
    [
        pytest.param(
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': {'value': 'IN DYNAMO'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': 'REF:initial/data'},
                {'PK': 'organization#foo', 'SK': 'keyspace#yet_another_key', 'data': 'REF:initial/data2'},
            ],
            {'initial/data': 'INITIAL', 'initial/data2': 'INITIAL2'},
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#other_key'},
            ],
            [{'value': 'IN DYNAMO'}, {'value': 'INITIAL2'}],
            id='delete_single_reference',
        ),
        pytest.param(
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': 'REF:initial/data'},
                {'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': {'value': 'IN DYNAMO'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#yet_another_key', 'data': 'REF:initial/data2'},
            ],
            {'initial/data': 'INITIAL'},
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#other_key'},
                {'PK': 'organization#foo', 'SK': 'keyspace#yet_another_key'},
            ],
            [{'value': 'IN DYNAMO'}],
            id='delete_all_references',
        ),
    ],
)
def test_list_objects_should_not_return_deleted_s3_references(
    context_dynamo, initial_dynamo_data, s3_data, removed_dynamo_data, expected_list
):
    for path, value in s3_data.items():
        MockEx(context_dynamo.mock_s3.get_json_object).if_called_with(
            context_dynamo.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}'
        ).return_value({'value': value})

    context_dynamo.load_data(initial_dynamo_data)

    # get the list generator
    gen = feature_data.list_objects(context_dynamo.conn, 'foo', 'keyspace')

    # change the dynamo reference and s3 contents
    for key in removed_dynamo_data:
        context_dynamo.conn._table.delete_item(Key=key)

    # resolving the generator to a list should get the updated data
    assert list(gen) == expected_list


@pytest.mark.unit
def test_list_objects_should_fail_if_key_is_still_missing(context_dynamo):
    dynamo_data = [
        {'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': {'value': 'IN DYNAMO'}},
        {'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': 'REF:path/to/initial/data'},
    ]
    context_dynamo.load_data(dynamo_data)

    # the object is completely missing from S3
    MockEx(context_dynamo.mock_s3.get_json_object).if_called_with(
        context_dynamo.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/path/to/initial/data'
    ).raises(NOT_FOUND_ERROR)

    spy = context_dynamo.mocker.spy(context_dynamo.conn._table, 'get_item')
    with pytest.raises(ClientError):
        list(feature_data.list_objects(context_dynamo.conn, 'foo', 'keyspace'))

    assert len(spy.mock_calls) == MAX_RESOLVE_DATA_TRIES


@pytest.mark.unit
def test_list_objects_should_fail_if_other_error_happens(context_dynamo):
    dynamo_data = [
        {'PK': 'organization#foo', 'SK': 'keyspace#some_key', 'data': {'value': 'IN DYNAMO'}},
        {'PK': 'organization#foo', 'SK': 'keyspace#other_key', 'data': 'REF:path/to/initial/data'},
    ]
    context_dynamo.load_data(dynamo_data)

    # some other error happens when getting the data
    MockEx(context_dynamo.mock_s3.get_json_object).if_called_with(
        context_dynamo.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/path/to/initial/data'
    ).raises(Exception('SENTINEL'))

    spy = context_dynamo.mocker.spy(context_dynamo.conn._table, 'get_item')
    with pytest.raises(Exception, match='SENTINEL'):
        list(feature_data.list_objects(context_dynamo.conn, 'foo', 'keyspace'))

    assert len(spy.mock_calls) == 0


@pytest.mark.unit
@pytest.mark.parametrize(
    'keyspace,keys,expected_dynamo_condition',
    [
        ('ks', [], Attr('PK').begins_with('organization#') & Attr('SK').begins_with('ks#')),
        ('obj', ['k1', 'k2'], Attr('PK').begins_with('organization#') & Attr('SK').begins_with('obj#k1#k2#')),
    ],
)
def test_list_all_objects_should_use_correct_dynamo_condition(context, keyspace, keys, expected_dynamo_condition):
    list(feature_data.list_all_objects(context.conn, keyspace, *keys))

    context.mock_dyn.scan.assert_called_once_with(context.conn._table, dict(FilterExpression=expected_dynamo_condition))


@pytest.mark.unit
@pytest.mark.parametrize(
    'dynamo_data,expected_result',
    [
        ([], []),
        (
            [
                {'PK': 'organization#abc123', 'SK': 'keyspace#my_key', 'data': {'val': 1}},
                {'PK': 'organization#def456', 'SK': 'keyspace#other_key', 'data': {'val': 2}},
            ],
            [{'val': 1}, {'val': 2}],
        ),
    ],
)
def test_list_all_objects_should_return_dynamo_data(context, dynamo_data, expected_result):
    MockEx(context.mock_dyn.scan).if_called_with_anything().return_value(dynamo_data)
    result = list(feature_data.list_all_objects(context.conn, context.keyspace))

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'dynamo_data,expected_s3_data,expected_result',
    [
        (
            [{'PK': 'organization#abc123', 'SK': 'keyspace#my_key', 'data': 'REF:path/to/data'}],
            {'path/to/data': {'data': 1}},
            [{'data': 1}],
        ),
        (
            [
                {'PK': 'organization#abc123', 'SK': 'keyspace#my_key', 'data': 'REF:path/to/data'},
                {'PK': 'organization#def456', 'SK': 'keyspace#your_key', 'data': 'REF:path2'},
            ],
            {'path/to/data': {'data': 1}, 'path2': {'data': 2}},
            [{'data': 1}, {'data': 2}],
        ),
        (
            [
                {'PK': 'organization#abc123', 'SK': 'keyspace#my_key', 'data': 'REF:path/to/data'},
                {'PK': 'organization#def456', 'SK': 'keyspace#other_key', 'data': {'data': 2}},
            ],
            {'path/to/data': {'data': 1}},
            [{'data': 1}, {'data': 2}],
        ),
        ([{'PK': 'organization#baddata', 'SK': 'keyspace#my_key', 'data': 'Invalid string'}], {}, ValueError),
    ],
)
def test_list_all_objects_should_return_s3_data_for_ref(context_dynamo, dynamo_data, expected_s3_data, expected_result):
    context_dynamo.load_data(dynamo_data)

    for path, data in expected_s3_data.items():
        MockEx(context_dynamo.mock_s3.get_json_object).if_called_with(
            context_dynamo.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}'
        ).return_value(data)

    def list_objects():
        return list(feature_data.list_all_objects(context_dynamo.conn, 'keyspace'))

    if expected_result is ValueError:
        pytest.raises(ValueError, list_objects)
    else:
        assert list_objects() == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'keyspace,keys,expected_dynamo_condition',
    [
        ('ks', [], Attr('PK').begins_with('organization#') & Attr('SK').begins_with('ks#')),
        ('obj', ['k1', 'k2'], Attr('PK').begins_with('organization#') & Attr('SK').begins_with('obj#k1#k2#')),
    ],
)
def test_list_all_objects_with_organization_should_use_correct_dynamo_condition(
    context, keyspace, keys, expected_dynamo_condition
):
    list(feature_data.list_all_objects_with_organization(context.conn, keyspace, *keys))

    context.mock_dyn.scan.assert_called_once_with(context.conn._table, dict(FilterExpression=expected_dynamo_condition))


@pytest.mark.unit
@pytest.mark.parametrize(
    'dynamo_data,expected_result',
    [
        ([], []),
        (
            [
                {'PK': 'organization#abc123', 'SK': 'keyspace#my_key', 'data': {'val': 1}},
                {'PK': 'organization#def456', 'SK': 'keyspace#other_key', 'data': {'val': 2}},
            ],
            [('abc123', {'val': 1}), ('def456', {'val': 2})],
        ),
    ],
)
def test_list_all_objects_with_organization_should_return_dynamo_data(context, dynamo_data, expected_result):
    MockEx(context.mock_dyn.scan).if_called_with_anything().return_value(dynamo_data)
    result = list(feature_data.list_all_objects_with_organization(context.conn, context.keyspace))

    assert result == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'dynamo_data,mock_s3_data,expected_result',
    [
        (
            [{'PK': 'organization#abc123', 'SK': 'keyspace#my_key', 'data': 'REF:path/to/data'}],
            {'path/to/data': {'data': 1}},
            [('abc123', {'data': 1})],
        ),
        (
            [
                {'PK': 'organization#abc123', 'SK': 'keyspace#my_key', 'data': 'REF:path/to/data'},
                {'PK': 'organization#def456', 'SK': 'keyspace#my_other_key', 'data': 'REF:path2'},
            ],
            {'path/to/data': {'data': 1}, 'path2': {'data': 2}},
            [('abc123', {'data': 1}), ('def456', {'data': 2})],
        ),
        (
            [
                {'PK': 'organization#abc123', 'SK': 'keyspace#my_key', 'data': 'REF:path/to/data'},
                {'PK': 'organization#def456', 'SK': 'keyspace#your_key', 'data': {'data': 2}},
            ],
            {'path/to/data': {'data': 1}},
            [('abc123', {'data': 1}), ('def456', {'data': 2})],
        ),
        ([{'PK': 'organization#baddata', 'SK': 'keyspace#my_key', 'data': 'Invalid string'}], {}, ValueError),
    ],
)
def test_list_all_objects_with_organization_should_return_s3_data_for_ref(context, dynamo_data, mock_s3_data, expected_result):
    MockEx(context.mock_dyn.scan).if_called_with_anything().return_value(dynamo_data)
    for path, data in mock_s3_data.items():
        MockEx(context.mock_s3.get_json_object).if_called_with(
            context.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}'
        ).return_value(data)

    def list_objects():
        return list(feature_data.list_all_objects_with_organization(context.conn, context.keyspace))

    if expected_result is ValueError:
        pytest.raises(ValueError, list_objects)
    else:
        assert list_objects() == expected_result


@pytest.mark.unit
@pytest.mark.parametrize(
    'dynamo_data,initial_s3_data,updated_s3_data,updated_dynamo_data,expected',
    [
        pytest.param(
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key1', 'data': {'value': 'FOO: IN DYNAMO'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key2', 'data': {'value': 'FOO: IN DYNAMO2'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/foo'},
                {'PK': 'organization#bar', 'SK': 'keyspace#dyn_key', 'data': {'value': 'BAR: IN DYNAMO'}},
                {'PK': 'organization#bar', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/bar'},
            ],
            {'path/to/initial/foo': 'FOO: IN S3 INITIAL', 'path/to/initial/bar': 'BAR: IN S3 INITIAL'},
            {'path/to/updated/foo': 'FOO: IN S3 UPDATED', 'path/to/updated/bar': 'BAR: IN S3 UPDATED'},
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/updated/foo'},
                {'PK': 'organization#bar', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/updated/bar'},
            ],
            [
                ('foo', {'value': 'FOO: IN DYNAMO'}),
                ('foo', {'value': 'FOO: IN DYNAMO2'}),
                ('foo', {'value': 'FOO: IN S3 UPDATED'}),
                ('bar', {'value': 'BAR: IN DYNAMO'}),
                ('bar', {'value': 'BAR: IN S3 UPDATED'}),
            ],
            id='both_orgs_have_changed_data',
        ),
        pytest.param(
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key1', 'data': {'value': 'FOO: IN DYNAMO'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key2', 'data': {'value': 'FOO: IN DYNAMO2'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/foo'},
                {'PK': 'organization#bar', 'SK': 'keyspace#dyn_key', 'data': {'value': 'BAR: IN DYNAMO'}},
                {'PK': 'organization#bar', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/bar'},
            ],
            {'path/to/initial/foo': 'FOO: IN S3 INITIAL', 'path/to/initial/bar': 'BAR: IN S3 INITIAL'},
            {'path/to/initial/foo': 'FOO: IN S3 INITIAL', 'path/to/updated/bar': 'BAR: IN S3 UPDATED'},
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/foo'},
                {'PK': 'organization#bar', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/updated/bar'},
            ],
            [
                ('foo', {'value': 'FOO: IN DYNAMO'}),
                ('foo', {'value': 'FOO: IN DYNAMO2'}),
                ('foo', {'value': 'FOO: IN S3 INITIAL'}),
                ('bar', {'value': 'BAR: IN DYNAMO'}),
                ('bar', {'value': 'BAR: IN S3 UPDATED'}),
            ],
            id='only_one_org_changed',
        ),
        pytest.param(
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key1', 'data': {'value': 'FOO: IN DYNAMO'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key2', 'data': {'value': 'FOO: IN DYNAMO2'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/foo'},
                {'PK': 'organization#bar', 'SK': 'keyspace#dyn_key', 'data': {'value': 'BAR: IN DYNAMO'}},
                {'PK': 'organization#bar', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/bar'},
            ],
            {'path/to/initial/foo': 'FOO: IN S3 INITIAL', 'path/to/initial/bar': 'BAR: IN S3 INITIAL'},
            {},
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#s3_key', 'data': {'value': 'FOO: BACK IN DYNAMO'}},
                {'PK': 'organization#bar', 'SK': 'keyspace#s3_key', 'data': {'value': 'BAR: BACK IN DYNAMO'}},
            ],
            [
                ('foo', {'value': 'FOO: IN DYNAMO'}),
                ('foo', {'value': 'FOO: IN DYNAMO2'}),
                ('foo', {'value': 'FOO: BACK IN DYNAMO'}),
                ('bar', {'value': 'BAR: IN DYNAMO'}),
                ('bar', {'value': 'BAR: BACK IN DYNAMO'}),
            ],
            id='data_changed_from_ref_to_dynamo',
        ),
    ],
)
def test_list_all_objects_with_organization_should_reload_changed_s3_references(
    context_dynamo, dynamo_data, initial_s3_data, updated_s3_data, updated_dynamo_data, expected
):
    def mock_s3(remove: dict, add: dict):
        for path in remove.keys():
            MockEx(context_dynamo.mock_s3.get_json_object).if_called_with(
                context_dynamo.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}'
            ).raises(NOT_FOUND_ERROR)
        for path, value in add.items():
            MockEx(context_dynamo.mock_s3.get_json_object).if_called_with(
                context_dynamo.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{path}'
            ).return_value({'value': value})

    mock_s3({}, initial_s3_data)
    context_dynamo.load_data(dynamo_data)

    gen = feature_data.list_all_objects_with_organization(context_dynamo.conn, 'keyspace')

    mock_s3(initial_s3_data, updated_s3_data)
    context_dynamo.load_data(updated_dynamo_data)

    # resolving the generator to a list should get the updated data
    assert list(gen) == expected


@pytest.mark.unit
@pytest.mark.parametrize(
    'initial_dynamo_data,removed_dynamo_data,expected',
    [
        pytest.param(
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key1', 'data': {'value': 'FOO: IN DYNAMO'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key2', 'data': {'value': 'FOO: IN DYNAMO2'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/foo'},
                {'PK': 'organization#bar', 'SK': 'keyspace#dyn_key', 'data': {'value': 'BAR: IN DYNAMO'}},
                {'PK': 'organization#bar', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/bar'},
            ],
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#s3_key'},
            ],
            [
                ('foo', {'value': 'FOO: IN DYNAMO'}),
                ('foo', {'value': 'FOO: IN DYNAMO2'}),
                ('bar', {'value': 'BAR: IN DYNAMO'}),
                ('bar', {'value': 'FROM S3'}),
            ],
            id='remove_one_org',
        ),
        pytest.param(
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key1', 'data': {'value': 'FOO: IN DYNAMO'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key2', 'data': {'value': 'FOO: IN DYNAMO2'}},
                {'PK': 'organization#foo', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/foo'},
                {'PK': 'organization#bar', 'SK': 'keyspace#dyn_key', 'data': {'value': 'BAR: IN DYNAMO'}},
                {'PK': 'organization#bar', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/bar'},
            ],
            [
                {'PK': 'organization#foo', 'SK': 'keyspace#s3_key'},
                {'PK': 'organization#bar', 'SK': 'keyspace#s3_key'},
            ],
            [
                ('foo', {'value': 'FOO: IN DYNAMO'}),
                ('foo', {'value': 'FOO: IN DYNAMO2'}),
                ('bar', {'value': 'BAR: IN DYNAMO'}),
            ],
            id='remove_both_orgs',
        ),
    ],
)
def test_list_all_objects_with_deleted_keys(context_dynamo, initial_dynamo_data, removed_dynamo_data, expected):
    context_dynamo.load_data(initial_dynamo_data)
    MockEx(context_dynamo.mock_s3.get_json_object).if_called_with_anything().return_value({'value': 'FROM S3'})

    gen = feature_data.list_all_objects_with_organization(context_dynamo.conn, 'keyspace')

    for key in removed_dynamo_data:
        context_dynamo.conn._table.delete_item(Key=key)

    # resolving the generator to a list should get the updated data
    assert list(gen) == expected


@pytest.mark.unit
def test_list_all_objects_with_organization_should_fail_if_key_is_still_missing(context_dynamo):
    dynamo_data = [
        {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key', 'data': {'value': 'foo: IN DYNAMO'}},
        {'PK': 'organization#foo', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/foo'},
        {'PK': 'organization#bar', 'SK': 'keyspace#dyn_key', 'data': {'value': 'bar: IN DYNAMO'}},
        {'PK': 'organization#bar', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/bar'},
    ]
    context_dynamo.load_data(dynamo_data)

    # the object is completely missing from S3 for "bar" org
    MockEx(context_dynamo.mock_s3.get_json_object).if_called_with(
        context_dynamo.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/path/to/initial/bar'
    ).raises(NOT_FOUND_ERROR)

    spy = context_dynamo.mocker.spy(context_dynamo.conn._table, 'get_item')
    with pytest.raises(ClientError):
        list(feature_data.list_all_objects_with_organization(context_dynamo.conn, 'keyspace'))

    assert len(spy.mock_calls) == MAX_RESOLVE_DATA_TRIES


@pytest.mark.unit
def test_list_all_objects_with_organization_should_fail_if_other_error_happens(context_dynamo):
    dynamo_data = [
        {'PK': 'organization#foo', 'SK': 'keyspace#dyn_key', 'data': {'value': 'foo: IN DYNAMO'}},
        {'PK': 'organization#foo', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/foo'},
        {'PK': 'organization#bar', 'SK': 'keyspace#dyn_key', 'data': {'value': 'bar: IN DYNAMO'}},
        {'PK': 'organization#bar', 'SK': 'keyspace#s3_key', 'data': 'REF:path/to/initial/bar'},
    ]
    context_dynamo.load_data(dynamo_data)

    # some other error happens when getting the data
    MockEx(context_dynamo.mock_s3.get_json_object).if_called_with(
        context_dynamo.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/path/to/initial/foo'
    ).raises(Exception('SENTINEL'))
    spy = context_dynamo.mocker.spy(context_dynamo.conn._table, 'get_item')

    with pytest.raises(Exception, match='SENTINEL'):
        list(feature_data.list_all_objects_with_organization(context_dynamo.conn, 'keyspace'))

    assert len(spy.mock_calls) == 0


@pytest.mark.unit
@pytest.mark.parametrize(
    'org_id,keyspace,keys,expected_dynamo_key',
    [
        ('org1', 'ks', [], dict(PK='organization#org1', SK='ks')),
        ('org2', 'obj', ['k1', 'k2'], dict(PK='organization#org2', SK='obj#k1#k2')),
    ],
)
def test_set_object_should_use_correct_dynamo_key(context, org_id, keyspace, keys, expected_dynamo_key):
    feature_data.set_object(context.conn, org_id, keyspace, *keys, data={})

    context.mock_dyn.put_item.assert_called_once_with(
        context.conn._table, ReturnValues='ALL_OLD', Item=dict(**expected_dynamo_key, data=ANY)
    )


@pytest.mark.unit
@pytest.mark.parametrize(
    'data,expected_put_data', [({}, {}), ({'val': 1}, {'val': 1}), ({'val': 1.23}, {'val': Decimal('1.23')})]
)
def test_set_object_should_put_dynamo_data(context, data, expected_put_data):
    feature_data.set_object(context.conn, context.org_id, context.keyspace, data=data)

    context.mock_dyn.put_item.assert_called_once_with(
        context.conn._table, ReturnValues='ALL_OLD', Item=dict(PK=ANY, SK=ANY, data=expected_put_data)
    )


@pytest.mark.unit
@pytest.mark.parametrize(
    'org_id,keyspace,keys,expected_ref',
    [
        ('org1', 'ks', [], 'organization_org1/ks_UUID.json'),
        ('aaa-bbb', 'obj', ['key@1', 'key@2'], 'organization_aaa_bbb/obj_key_1_key_2_UUID.json'),
    ],
)
def test_set_object_should_put_large_data_in_s3_with_ref(context, org_id, keyspace, keys, expected_ref):
    context.mock_dyn.put_item.side_effect = [EXCEED_MAX_SIZE_ERROR, None]
    feature_data.set_object(context.conn, org_id, keyspace, *keys, data={})

    context.mock_dyn.put_item.assert_called_with(
        context.conn._table, ReturnValues='ALL_OLD', Item=dict(PK=ANY, SK=ANY, data=f'REF:{expected_ref}')
    )
    context.mock_s3.put_json_object.assert_called_once_with(
        context.conn._s3_bucket, f'{S3_OBJECT_STORE_ROOT}/{expected_ref}', ANY
    )


@pytest.mark.unit
@pytest.mark.parametrize('data,expected_s3_data', [({}, {}), ({'val': 1}, {'val': 1})])
def test_set_object_should_put_large_data_in_s3(context, data, expected_s3_data):
    context.mock_dyn.put_item.side_effect = [EXCEED_MAX_SIZE_ERROR, None]
    feature_data.set_object(context.conn, context.org_id, context.keyspace, data=data)

    context.mock_s3.put_json_object.assert_called_once_with(context.conn._s3_bucket, ANY, expected_s3_data)


@pytest.mark.parametrize('prior_data,expected_delete_key', [({'val': 1}, None), ('REF:key.json', 'key.json')])
def test_set_object_should_delete_prior_data_from_s3(context, prior_data, expected_delete_key):
    context.mock_dyn.put_item.return_value = dict(Attributes=dict(data=prior_data))
    feature_data.set_object(context.conn, context.org_id, context.keyspace, data={})

    if expected_delete_key is not None:
        context.mock_s3_client.delete_object.assert_called_once_with(
            Bucket=context.conn._s3_bucket, Key=f'{S3_OBJECT_STORE_ROOT}/{expected_delete_key}'
        )
    else:
        assert not context.mock_s3_client.delete_object.called


@pytest.mark.unit
@pytest.mark.parametrize(
    'current_data,expected_object',
    [
        ({}, {}),
        ({'data': None}, {}),
        ({'data': {}}, {}),
        ({'data': {'key': 'value'}}, {'key': 'value'}),
        ({'data': 'BAD DATA'}, ValueError),
    ],
)
def test_update_object_data_atomic_should_call_handler_with_current_data(context, current_data, expected_object):
    MockEx(context.conn._table.get_item).if_called_with_anything().return_value({'Item': current_data})

    mock_update_function = Mock()

    def update():
        feature_data.update_object_atomic(context.conn, context.org_id, context.keyspace, update_handler=mock_update_function)

    if expected_object is ValueError:
        pytest.raises(expected_object, update)
    else:
        update()
        mock_update_function.assert_called_once_with(expected_object)


@pytest.mark.unit
@pytest.mark.parametrize(
    'new_data,should_put_item', [(None, False), ({'key': 'value'}, True), ({'key': 'value', 'k2': 123}, True)]
)
def test_update_object_data_atomic_should_put_new_data_with_handler_result(context, new_data, should_put_item):
    mock_update_function = Mock()
    mock_update_function.return_value = new_data

    feature_data.update_object_atomic(context.conn, context.org_id, context.keyspace, update_handler=mock_update_function)

    assert context.mock_dyn.put_item.called == should_put_item
    if should_put_item:
        context.mock_dyn.put_item.assert_called_with(
            context.conn._table, ConditionExpression=ANY, ReturnValues='ALL_OLD', Item=dict(PK=ANY, SK=ANY, data=new_data)
        )


@pytest.mark.unit
@pytest.mark.parametrize(
    'exception,should_retry',
    [
        (ClientError(error_response=dict(Error=dict(Code='ConditionalCheckFailedException')), operation_name='PutItem'), True),
        (ValueError('other error'), False),
        (ClientError(error_response=dict(Error=dict(Code='ValidationException')), operation_name='PutItem'), False),
    ],
)
def test_update_object_data_atomic_should_retry_if_condition_check_fails(context, exception, should_retry):
    new_data = {'test': 'data'}
    mock_update_function = Mock()
    mock_update_function.return_value = new_data
    context.mock_dyn.put_item.side_effect = [exception, None]

    def update():
        feature_data.update_object_atomic(context.conn, context.org_id, context.keyspace, update_handler=mock_update_function)

    expected_call = call(
        context.conn._table, ConditionExpression=ANY, ReturnValues='ALL_OLD', Item=dict(PK=ANY, SK=ANY, data=new_data)
    )
    if should_retry:
        update()
        assert context.mock_dyn.put_item.mock_calls == [expected_call, expected_call]
    else:
        pytest.raises(type(exception), update)


@pytest.mark.unit
@pytest.mark.parametrize(
    'org_id,keyspace,keys,expected_dynamo_key',
    [
        ('org1', 'ks', [], dict(PK='organization#org1', SK='ks')),
        ('org2', 'obj', ['k1', 'k2'], dict(PK='organization#org2', SK='obj#k1#k2')),
    ],
)
def test_delete_object_should_delete_dynamo_data(context, org_id, keyspace, keys, expected_dynamo_key):
    feature_data.delete_object(context.conn, org_id, keyspace, *keys)

    context.conn._table.delete_item.assert_called_once_with(ReturnValues='ALL_OLD', Key=expected_dynamo_key)


@pytest.mark.unit
@pytest.mark.parametrize('prior_data,expected_delete_key', [({'val': 1}, None), ('REF:key.json', 'key.json')])
def test_delete_object_should_delete_prior_data_from_s3(context, prior_data, expected_delete_key):
    context.conn._table.delete_item.return_value = dict(Attributes=dict(data=prior_data))
    feature_data.delete_object(context.conn, context.org_id, context.keyspace)

    if expected_delete_key is not None:
        context.mock_s3_client.delete_object.assert_called_once_with(
            Bucket=context.conn._s3_bucket, Key=f'{S3_OBJECT_STORE_ROOT}/{expected_delete_key}'
        )
    else:
        assert not context.mock_s3_client.delete_object.called
