Faster China - tests (#169)

Author: nirLumigo

src/lumigo_tracer/lumigo_utils.py

@@ -371,9 +371,6 @@ def _send_data_to_kinesis(
     if not boto3:
         get_logger().error("boto3 is missing. Unable to send to Kinesis.")
         return None
-    if not botocore:
-        get_logger().error("botocore is missing. Unable to send to Kinesis.")
-        return None
     client = _get_edge_kinesis_boto_client(
         region=region,
         aws_access_key_id=aws_access_key_id,

src/test/conftest.py

@@ -28,6 +28,7 @@ def cancel_timeout_mechanism(monkeypatch):
 @pytest.fixture(autouse=True)
 def remove_caches(monkeypatch):
     get_omitting_regex.cache_clear()
+    monkeypatch.setattr(lumigo_utils, "edge_kinesis_boto_client", None)
 
 
 @pytest.yield_fixture(autouse=True)

src/test/unit/test_lumigo_utils.py

@@ -1,9 +1,13 @@
+import importlib.util
 import inspect
+import logging
 from collections import OrderedDict
 from decimal import Decimal
 import datetime
 import http.client
-from mock import Mock
+
+import boto3
+from mock import Mock, MagicMock
 
 import pytest
 from lumigo_tracer import lumigo_utils
@@ -427,6 +431,53 @@ def test_report_json_china_no_boto(monkeypatch, reporter_mock, caplog):
     )
 
 
+def test_report_json_china_on_error_no_exception_and_notify_user(capsys, monkeypatch):
+    monkeypatch.setattr(Configuration, "should_report", True)
+    monkeypatch.setattr(Configuration, "edge_kinesis_aws_access_key_id", "my_value")
+    monkeypatch.setattr(Configuration, "edge_kinesis_aws_secret_access_key", "my_value")
+    monkeypatch.setattr(boto3, "client", MagicMock(side_effect=Exception))
+    lumigo_utils.get_logger().setLevel(logging.CRITICAL)
+
+    report_json(CHINA_REGION, [{"a": "b"}])
+
+    assert "Failed to send spans" in capsys.readouterr().out
+
+
+def test_china_shouldnt_establish_http_connection(monkeypatch):
+    monkeypatch.setenv("AWS_REGION", CHINA_REGION)
+    # Reload a duplicate of lumigo_utils
+    spec = importlib.util.find_spec("lumigo_tracer.lumigo_utils")
+    lumigo_utils_reloaded = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(lumigo_utils_reloaded)
+
+    assert lumigo_utils_reloaded.edge_connection is None
+
+
+def test_china_with_env_variable_shouldnt_reuse_boto3_connection(monkeypatch):
+    monkeypatch.setenv("LUMIGO_KINESIS_SHOULD_REUSE_CONNECTION", "false")
+    monkeypatch.setattr(Configuration, "should_report", True)
+    monkeypatch.setattr(Configuration, "edge_kinesis_aws_access_key_id", "my_value")
+    monkeypatch.setattr(Configuration, "edge_kinesis_aws_secret_access_key", "my_value")
+    monkeypatch.setattr(boto3, "client", MagicMock())
+
+    report_json(CHINA_REGION, [{"a": "b"}])
+    report_json(CHINA_REGION, [{"a": "b"}])
+
+    assert boto3.client.call_count == 2
+
+
+def test_china_reuse_boto3_connection(monkeypatch):
+    monkeypatch.setattr(Configuration, "should_report", True)
+    monkeypatch.setattr(Configuration, "edge_kinesis_aws_access_key_id", "my_value")
+    monkeypatch.setattr(Configuration, "edge_kinesis_aws_secret_access_key", "my_value")
+    monkeypatch.setattr(boto3, "client", MagicMock())
+
+    report_json(CHINA_REGION, [{"a": "b"}])
+    report_json(CHINA_REGION, [{"a": "b"}])
+
+    boto3.client.assert_called_once()
+
+
 @pytest.mark.parametrize("env, expected", [("True", True), ("other", False), ("123", False)])
 def test_is_kill_switch_on(monkeypatch, env, expected):
     monkeypatch.setenv(KILL_SWITCH, env)

src/test/unit/test_tracer.py

@@ -6,7 +6,7 @@
 import sys
 from functools import wraps
 import logging
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, ANY
 
 import boto3
 import pytest
@@ -419,9 +419,10 @@ def lambda_test_function(event, context):
     span_sent = json.loads(records[1]["Data"].decode())[0]
     assert span_sent["event"] == json.dumps(event)
     # Used the client from the decorator params
-    # boto3.client.assert_called_with(
-    #     "kinesis",
-    #     region_name=china_region_for_test,
-    #     aws_access_key_id=access_key_id,
-    #     aws_secret_access_key=secret_access_key,
-    # )
+    boto3.client.assert_called_with(
+        "kinesis",
+        region_name=china_region_for_test,
+        aws_access_key_id=access_key_id,
+        aws_secret_access_key=secret_access_key,
+        config=ANY,
+    )