Aws lambda unit test throwing SegmentNotFoundException

I'm trying to write unit tests for an AWS Lambda function that uses the boto3 library to interact with AWS S3. I'm using the moto library to mock S3 in my tests. However, I'm running into an error that says raise SegmentNotFoundException(MISSING_SEGMENT_MSG) aws_xray_sdk.core.exceptions.exceptions.SegmentNotFoundException: cannot find the current segment/subsegment, please make sure you have a segment open

I've tried using @mock.patch('boto3.client') to mock the S3 client and @mock_s3 to mock the S3 service, but I'm still getting the same error.

Here's a simplified version of my test code:

import unittest
from moto import mock_s3
import boto3

class TestMyLambda(unittest.TestCase):
@mock_s3
def test_my_lambda(self):
bucket_name = 'my-test-bucket'
s3_client = boto3.client('s3', region_name='us-east-1')
s3_client.create_bucket(Bucket=bucket_name)
\# ... rest of the test code ...

if __name__ == '__main__':
unittest.main()`

In my lambda handler I am capturing the x-ray traces:

@xray_recorder.capture('my-lambda')
@logger.inject_lambda_context()
def lambda_handler(event, context):
    """
    Handler will checks if the json paylod is same as the one saved in s3. If there is a change the the file is replaced with latest entries, else ignored
    """
    xray_recorder.begin_segment('extension-process')

    # Remaining code .......

I'm not sure why I'm getting this error and how to fix it. Can someone help me understand what's going on here and how to properly mock S3 in my tests?

X-Ray is auto-instrumenting the code and this fails in local test setups. You either need to setup Xray locally so it can also run in unit tests or disable it for unit testing.

To disable it set the environment variable AWS_XRAY_SDK_ENABLED to false. Then X-Ray will generate dummy-segments and the error will go away.

According to documentation, it should also be possible to set AWS_XRAY_CONTEXT_MISSING to IGNORE_ERROR to not throw exceptions, but I can't try it right now.