def construct_ctx():
"""
Clean up context storage on each test run and begin a segment
so that later subsegment can be attached. After each test run
it cleans up context storage again.
"""
pre_run_modules = set(module for module in sys.modules.keys())
xray_recorder.configure(service='test', sampling=False, context=Context())
xray_recorder.clear_trace_entities()
xray_recorder.begin_segment('name')
yield
xray_recorder.end_segment()
xray_recorder.clear_trace_entities()
global_sdk_config.set_sdk_enabled(True)
# Reload wrapt.importer references to modules to start off clean
reload(wrapt)
reload(wrapt.importer)
# Reload patcher references to already patched modules
reload(patcher)
# Cleanup the already imported module references in the system
for module_name, module in sorted(sys.modules.items(),
key=lambda m: len(m[0]),
reverse=True):
if module_name not in pre_run_modules and inspect.ismodule(module):
reload(module)
for module_name in sorted(sys.modules.keys(),
key=lambda m: len(m),
reverse=True):
if module_name not in pre_run_modules:
del sys.modules[module_name]
def test_disabled_force_sampling():
global_sdk_config.set_sdk_enabled(False)
xray_recorder.configure(sampling=True)
segment_entity = xray_recorder.begin_segment("Test1")
subsegment_entity = xray_recorder.begin_subsegment("Test2")
assert type(segment_entity) is DummySegment
assert type(subsegment_entity) is DummySubsegment
def cleanup():
"""
Clean up context storage before and after each test run
"""
recorder.clear_trace_entities()
yield
recorder.clear_trace_entities()
global_sdk_config.set_sdk_enabled(True)
def test_disable_sdk_disables_patching():
global_sdk_config.set_sdk_enabled(False)
patcher.patch(['tests.mock_module'])
imported_modules = [
module for module in TEST_MODULES if module in sys.modules
]
assert not imported_modules
assert len(xray_recorder.current_segment().subsegments) == 0
def test_disabled_ends():
global_sdk_config.set_sdk_enabled(False)
xray_recorder.end_segment()
xray_recorder.end_subsegment()
xray_recorder.end_segment()
xray_recorder.end_segment()
xray_recorder.end_subsegment()
xray_recorder.end_subsegment()
def test_disabled_out_of_order_begins():
global_sdk_config.set_sdk_enabled(False)
xray_recorder.begin_subsegment("Test")
xray_recorder.begin_segment("Test")
xray_recorder.begin_subsegment("Test1")
xray_recorder.begin_subsegment("Test2")
assert type(xray_recorder.begin_subsegment("Test3")) is DummySubsegment
assert type(xray_recorder.begin_segment("Test4")) is DummySegment
def test_disable():
context.clear_trace_entities()
segment = context.get_trace_entity()
assert segment.sampled
context.clear_trace_entities()
global_sdk_config.set_sdk_enabled(False)
segment = context.get_trace_entity()
assert not segment.sampled
def construct_ctx():
"""
Clean up context storage before and after each test run.
"""
xray_recorder.configure(sampling=False)
xray_recorder.clear_trace_entities()
yield
xray_recorder.clear_trace_entities()
global_sdk_config.set_sdk_enabled(True)
def construct_ctx(monkeypatch):
"""
Clean up context storage before and after each test run.
"""
monkeypatch.delattr("botocore.session.Session.get_credentials")
xray_recorder.configure(sampling=False)
xray_recorder.clear_trace_entities()
yield
xray_recorder.clear_trace_entities()
global_sdk_config.set_sdk_enabled(True)
def cleanup():
"""
Clean up Environmental Variable for enable before and after tests
"""
if XRAY_ENABLED_KEY in os.environ:
del os.environ[XRAY_ENABLED_KEY]
yield
if XRAY_ENABLED_KEY in os.environ:
del os.environ[XRAY_ENABLED_KEY]
global_sdk_config.set_sdk_enabled(True)
def aws_credentials():
# Mocked AWS Credentials for moto, this is required
# as per their readme: https://github.com/spulec/moto
os.environ['AWS_ACCESS_KEY_ID'] = 'testing'
os.environ['AWS_SECRET_ACCESS_KEY'] = 'testing'
os.environ['AWS_SECURITY_TOKEN'] = 'testing'
os.environ['AWS_SESSION_TOKEN'] = 'testing'
# Disable X-Ray for unit tests
from aws_xray_sdk import global_sdk_config
global_sdk_config.set_sdk_enabled(False)
def init_xray(app):
logging.getLogger("aws_xray_sdk").setLevel('ERROR')
if app.debug:
# env var AWS_XRAY_SDK_ENABLED can overwrite this
global_sdk_config.set_sdk_enabled(False)
else:
# TODO: configure x-ray service
xray_recorder.configure(service="btbapi")
# Setup X-Ray Flask Integration
XRayMiddleware(app, xray_recorder)
# Setup X-Ray psycopg2, boto3 (aws sdk) Integration
patch(["psycopg2", "boto3"])
def test_disabled_sdk(client, recorder):
"""
Test a normal response when the SDK is disabled.
:param test_client: AioHttp test client fixture
:param loop: Eventloop fixture
:param recorder: X-Ray recorder fixture
"""
global_sdk_config.set_sdk_enabled(False)
resp = client.get("/")
assert resp.status_code == 200
segment = recorder.emitter.pop()
assert not segment
def recorder():
"""
Clean up context storage before and after each test run
"""
loop = asyncio.get_event_loop()
xray_recorder = get_new_stubbed_recorder()
xray_recorder.configure(service="test",
sampling=False,
context=AsyncContext(loop=loop))
xray_recorder.clear_trace_entities()
yield xray_recorder
global_sdk_config.set_sdk_enabled(True)
xray_recorder.clear_trace_entities()
async def test_disabled_sdk(test_client, loop, recorder):
"""
Test a normal response when the SDK is disabled.
:param test_client: AioHttp test client fixture
:param loop: Eventloop fixture
:param recorder: X-Ray recorder fixture
"""
global_sdk_config.set_sdk_enabled(False)
client = await test_client(ServerTest.app(loop=loop))
resp = await client.get('/')
assert resp.status == 200
segment = recorder.emitter.pop()
assert not segment
def test_env_var_precedence():
os.environ[XRAY_ENABLED_KEY] = "true"
# Env Variable takes precedence. This is called to activate the internal check
global_sdk_config.set_sdk_enabled(False)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "false"
global_sdk_config.set_sdk_enabled(False)
assert global_sdk_config.sdk_enabled() is False
os.environ[XRAY_ENABLED_KEY] = "false"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is False
os.environ[XRAY_ENABLED_KEY] = "true"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "true"
global_sdk_config.set_sdk_enabled(None)
assert global_sdk_config.sdk_enabled() is True
def recorder(loop):
"""
Clean up context storage before and after each test run
"""
xray_recorder = get_new_stubbed_recorder()
xray_recorder.configure(service='test',
sampling=False,
context=AsyncContext(loop=loop))
patcher = patch('aws_xray_sdk.ext.aiohttp.middleware.xray_recorder',
xray_recorder)
patcher.start()
xray_recorder.clear_trace_entities()
yield xray_recorder
global_sdk_config.set_sdk_enabled(True)
xray_recorder.clear_trace_entities()
patcher.stop()
def init_app(cls,
app: Insanic,
recorder: AsyncAWSXRayRecorder = None) -> None:
"""
Initializes Insanic to use Incendiary.
- This loads all default Incendiary configs.
- Validates connection information to X-Ray Daemon.
- Configures X-Ray SDK Recorder
- Attaches middlewares to start stop segments
- Replaces :code:`Service` object with :code:`IncendiaryService`
to trace interservice communications.
- Replaces asyncio task factory.
- Patches configured modules.
:param app: Your Insanic application/
:param recorder: If you want to use your own recorder.
"""
# checks to see if tracing can be enabled
cls.app = app
cls.load_config(app.config)
messages = cls._check_prerequisites(app)
if len(messages) == 0:
global_sdk_config.set_sdk_enabled(True)
app.xray_recorder = recorder or xray_recorder
cls.setup_middlewares(app)
cls.setup_client(app)
cls.setup_listeners(app)
patch(app.config.INCENDIARY_XRAY_PATCH_MODULES, raise_errors=False)
app.plugin_initialized("incendiary", cls)
else:
cls._handle_error(app, messages)
app.config.INCENDIARY_XRAY_ENABLED = False
global_sdk_config.set_sdk_enabled(False)
def test_invalid_env_string():
os.environ[XRAY_ENABLED_KEY] = "INVALID"
# Env Variable takes precedence. This is called to activate the internal check
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "1.0"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "1-.0"
global_sdk_config.set_sdk_enabled(False)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "T RUE"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
def test_disabled_put_methods():
global_sdk_config.set_sdk_enabled(False)
xray_recorder.put_annotation("Test", "Value")
xray_recorder.put_metadata("Test", "Value", "Namespace")
def setup():
yield
global_sdk_config.set_sdk_enabled(True)
def test_disabled_sdk():
global_sdk_config.set_sdk_enabled(False)
path = '/ok'
app.get(path)
segment = recorder.emitter.pop()
assert not segment
def test_env_enable_case():
os.environ[XRAY_ENABLED_KEY] = "TrUE"
# Env Variable takes precedence. This is called to activate the internal check
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "true"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "1"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "y"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "t"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "False"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is False
os.environ[XRAY_ENABLED_KEY] = "falSE"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is False
os.environ[XRAY_ENABLED_KEY] = "0"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is False
def setUp(self):
django.setup()
xray_recorder.configure(context=Context(), context_missing='LOG_ERROR')
xray_recorder.clear_trace_entities()
global_sdk_config.set_sdk_enabled(True)
r = requests.post(url)
logging.info("stopped status {}".format(r))
if r.status_code != 200:
logging.info("failed stopping the lambda : {}".format(r.json()))
else:
# TODO: at some point we need to pass any information that requests body/json throws
logging.info("lambda successfully stopped")
if __name__ == "__main__":
try:
# try_verify_credentials()
if IS_XRAY_ENABLE == "1":
global_sdk_config.set_sdk_enabled(True)
xray_recorder.configure(
service='ecs',
context_missing='LOG_ERROR',
daemon_address='xray-service.kube-system:2000',
plugins=('EC2Plugin', 'ECSPlugin'))
libs_to_patch = ('boto3', 'requests')
patch(libs_to_patch)
else:
global_sdk_config.set_sdk_enabled(False)
event_loop()
except ClientError as e:
errlog.log("ClientError Agent Event Loop {} [{}] POD:{}".format(
e.response['Error']['Code'], traceback.format_exc(), SELF_ID))
def test_disabled_begin_subsegment():
global_sdk_config.set_sdk_enabled(False)
subsegment_entity = xray_recorder.begin_subsegment("Test")
assert type(subsegment_entity) is DummySubsegment
def test_disable_is_dummy():
global_sdk_config.set_sdk_enabled(False)
segment = xray_recorder.begin_segment('name')
subsegment = xray_recorder.begin_subsegment('name')
assert type(xray_recorder.current_segment()) is DummySegment
assert type(xray_recorder.current_subsegment()) is DummySubsegment
def test_disabled_empty_context_current_calls():
global_sdk_config.set_sdk_enabled(False)
assert type(xray_recorder.current_segment()) is DummySegment
assert type(xray_recorder.current_subsegment()) is DummySubsegment
def test_disabled_get_context_entity():
global_sdk_config.set_sdk_enabled(False)
entity = xray_recorder.get_trace_entity()
assert type(entity) is DummySegment
def __disable_tracing_provider(self):
"""Forcefully disables tracing and patching"""
from aws_xray_sdk import global_sdk_config
global_sdk_config.set_sdk_enabled(False)