|
| 1 | +import sys |
| 2 | +import os |
| 3 | +from unittest import mock |
| 4 | +import pytest |
| 5 | + |
| 6 | +# Add the 'umap_narrative' directory to the Python path to allow the script to be imported |
| 7 | +umap_narrative_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'umap_narrative')) |
| 8 | +if umap_narrative_dir not in sys.path: |
| 9 | + sys.path.insert(0, umap_narrative_dir) |
| 10 | + |
| 11 | +# Now we can import the main function from the script to be tested |
| 12 | +from reset_conversation import main as reset_conversation_main |
| 13 | + |
| 14 | +@pytest.fixture |
| 15 | +def mock_boto3_resources(): |
| 16 | + """Mocks the get_boto_resource function to return mock objects for S3 and DynamoDB.""" |
| 17 | + |
| 18 | + # --- Create Mock for S3 --- |
| 19 | + mock_s3_resource = mock.MagicMock() |
| 20 | + mock_bucket = mock.MagicMock() |
| 21 | + mock_s3_object = mock.MagicMock() |
| 22 | + mock_s3_object.key = 'visualizations/test-rid/some_file.html' |
| 23 | + # Configure the mock bucket's objects.filter to return a list containing our mock object |
| 24 | + mock_bucket.objects.filter.return_value = [mock_s3_object] |
| 25 | + # The Bucket() method of the S3 resource will return our mock bucket |
| 26 | + mock_s3_resource.Bucket.return_value = mock_bucket |
| 27 | + |
| 28 | + # --- Create Mock for DynamoDB --- |
| 29 | + mock_dynamodb_resource = mock.MagicMock() |
| 30 | + mock_table = mock.MagicMock() |
| 31 | + # Configure the query/scan methods to return one fake item to trigger the deletion logic |
| 32 | + mock_table.query.return_value = {'Items': [{'pk': 'some_key', 'sk': 'some_sort_key'}]} |
| 33 | + mock_table.scan.return_value = {'Items': [{'pk': 'some_key', 'sk': 'some_sort_key'}]} |
| 34 | + # The Table() method of the DynamoDB resource will return our mock table |
| 35 | + mock_dynamodb_resource.Table.return_value = mock_table |
| 36 | + |
| 37 | + # --- Create the main mock that replaces get_boto_resource --- |
| 38 | + with mock.patch('reset_conversation.get_boto_resource') as mock_get_resource: |
| 39 | + # Define a side effect to return the correct mock based on service name |
| 40 | + def get_resource_side_effect(service_name): |
| 41 | + if service_name == 'dynamodb': |
| 42 | + return mock_dynamodb_resource |
| 43 | + if service_name == 's3': |
| 44 | + return mock_s3_resource |
| 45 | + return mock.MagicMock() |
| 46 | + |
| 47 | + mock_get_resource.side_effect = get_resource_side_effect |
| 48 | + |
| 49 | + # Yield the mocks to the test function |
| 50 | + yield { |
| 51 | + "get_resource": mock_get_resource, |
| 52 | + "dynamodb": mock_dynamodb_resource, |
| 53 | + "s3": mock_s3_resource, |
| 54 | + "table": mock_table, |
| 55 | + "bucket": mock_bucket |
| 56 | + } |
| 57 | + |
| 58 | +def test_reset_conversation_calls_all_services(mock_boto3_resources): |
| 59 | + """ |
| 60 | + Tests that the main reset script calls both DynamoDB and S3 deletion logic. |
| 61 | + """ |
| 62 | + test_zid = "12345" |
| 63 | + test_rid = "r_test_12345" |
| 64 | + |
| 65 | + # Run the main function with test arguments |
| 66 | + reset_conversation_main(zid=test_zid, rid=test_rid) |
| 67 | + |
| 68 | + # 1. Assert that our main mock was called for both services |
| 69 | + mock_boto3_resources["get_resource"].assert_any_call("dynamodb") |
| 70 | + mock_boto3_resources["get_resource"].assert_any_call("s3") |
| 71 | + |
| 72 | + # 2. Assert that the script tried to get a DynamoDB table |
| 73 | + # It will be called many times, so just check it was called at all |
| 74 | + mock_boto3_resources["dynamodb"].Table.assert_called() |
| 75 | + |
| 76 | + # 3. Assert that a deletion was attempted on a table |
| 77 | + # This confirms that the query/scan + delete loop was entered |
| 78 | + mock_table = mock_boto3_resources["table"] |
| 79 | + # Check that batch_writer (for query results) or delete_item (for single items) was called |
| 80 | + assert mock_table.batch_writer.called or mock_table.delete_item.called |
| 81 | + |
| 82 | + # 4. Assert that the script tried to access the S3 bucket |
| 83 | + mock_boto3_resources["s3"].Bucket.assert_called_with(mock.ANY) # bucket name is from env |
| 84 | + |
| 85 | + # 5. Assert that the script attempted to delete S3 objects |
| 86 | + mock_bucket = mock_boto3_resources["bucket"] |
| 87 | + mock_bucket.delete_objects.assert_called_once() |
| 88 | + |
| 89 | +def test_reset_conversation_skips_s3_if_no_rid(mock_boto3_resources): |
| 90 | + """ |
| 91 | + Tests that S3 deletion is skipped if no report_id (rid) is provided. |
| 92 | + """ |
| 93 | + test_zid = "54321" |
| 94 | + |
| 95 | + # Run the main function without the 'rid' argument |
| 96 | + reset_conversation_main(zid=test_zid, rid=None) |
| 97 | + |
| 98 | + # Assert that the DynamoDB deletion logic was still called |
| 99 | + mock_boto3_resources["get_resource"].assert_any_call("dynamodb") |
| 100 | + mock_boto3_resources["dynamodb"].Table.assert_called() |
| 101 | + |
| 102 | + # Assert that the S3 logic was SKIPPED |
| 103 | + mock_bucket = mock_boto3_resources["bucket"] |
| 104 | + mock_bucket.delete_objects.assert_not_called() |
| 105 | + |
0 commit comments