Merge custom resources on detection (#515)

Author: c24t

opencensus/common/monitored_resource/monitored_resource.py

@@ -43,16 +43,25 @@ def get_instance():
     :rtype: :class:`opencensus.common.resource.Resource` or None
     :return: A `Resource` configured for the current environment.
     """
+    resources = []
+    env_resource = resource.get_from_env()
+    if env_resource is not None:
+        resources.append(env_resource)
+
     if k8s_utils.is_k8s_environment():
-        return resource.Resource(_K8S_CONTAINER, k8s_utils.get_k8s_metadata())
+        resources.append(resource.Resource(
+            _K8S_CONTAINER, k8s_utils.get_k8s_metadata()))
+
     if is_gce_environment():
-        return resource.Resource(
+        resources.append(resource.Resource(
             _GCE_INSTANCE,
-            gcp_metadata_config.GcpMetadataConfig().get_gce_metadata())
-    if is_aws_environment():
-        return resource.Resource(
+            gcp_metadata_config.GcpMetadataConfig().get_gce_metadata()))
+    elif is_aws_environment():
+        resources.append(resource.Resource(
             _AWS_EC2_INSTANCE,
             (aws_identity_doc_utils.AwsIdentityDocumentUtils()
-             .get_aws_metadata()))
+             .get_aws_metadata())))
 
-    return None
+    if not resources:
+        return None
+    return resource.merge_resources(resources)

opencensus/common/resource/__init__.py

@@ -49,22 +49,31 @@
 _UNQUOTE_RE = re.compile(r'^([\'"]?)([^\1]*)(\1)$')
 
 
-def merge_resources(r1, r2):
-    """Merge two resources to get a new resource.
+def merge_resources(resource_list):
+    """Merge multiple resources to get a new resource.
 
-    :type r1: :class:`Resource`
-    :param r1: The first resource to merge, takes priority in conflicts.
+    Resources earlier in the list take precedence: if multiple resources share
+    a label key, use the value from the first resource in the list with that
+    key. The combined resource's type will be the first non-null type in the
+    list.
 
-    :type r2: :class:`Resource`
-    :param r2: The second resource to merge.
+    :type resource_list: list(:class:`Resource`)
+    :param resource_list: The list of resources to combine.
 
     :rtype: :class:`Resource`
     :return: The new combined resource.
     """
-    type_ = r1.type or r2.type
-    labels = copy(r2.labels)
-    labels.update(r1.labels)
-    return Resource(type_, labels)
+    if not resource_list:
+        raise ValueError
+    rtype = None
+    for rr in resource_list:
+        if rr.type:
+            rtype = rr.type
+            break
+    labels = {}
+    for rr in reversed(resource_list):
+        labels.update(rr.labels)
+    return Resource(rtype, labels)
 
 
 def check_ascii_256(string):
@@ -150,7 +159,7 @@ def merge(self, other):
         :rtype: :class:`Resource`
         :return: The new combined resource.
         """
-        return merge_resources(self, other)
+        return merge_resources([self, other])
 
 
 def unquote(string):

tests/unit/common/monitored_resource_util/test_monitored_resource.py

@@ -13,12 +13,18 @@
 # limitations under the License.
 
 from contextlib import contextmanager
-import mock
 import os
-import unittest
+import sys
 
 from opencensus.common.monitored_resource import monitored_resource
 
+if sys.version_info < (3,):
+    import unittest2 as unittest
+    import mock
+else:
+    import unittest
+    from unittest import mock
+
 
 @contextmanager
 def mock_mr_method(method, use):
@@ -39,50 +45,71 @@ def mock_use_aws(use):
     return mock_mr_method('is_aws_environment', use)
 
 
+def mock_oc_env():
+    return mock.patch.dict('os.environ', {
+        'OC_RESOURCE_TYPE': 'mock_resource_type',
+        'OC_RESOURCE_LABELS': 'mock_label_key=mock_label_value'
+    })
+
+
 @contextmanager
 def mock_k8s_env():
     with mock_use_k8s(True):
-        with mock_use_gce(False):
-            with mock_use_aws(False):
-                yield
+        yield
 
 
 @contextmanager
 def mock_gce_env():
-    with mock_use_k8s(False):
-        with mock_use_gce(True):
-            with mock_use_aws(False):
-                yield
+    with mock_use_gce(True):
+        with mock_use_aws(False):
+            yield
 
 
 @contextmanager
 def mock_aws_env():
-    with mock_use_k8s(False):
-        with mock_use_gce(False):
-            with mock_use_aws(True):
-                yield
+    with mock_use_gce(False):
+        with mock_use_aws(True):
+            yield
 
 
 class TestMonitoredResource(unittest.TestCase):
 
+    def setUp(self):
+        self.env_mock = mock.patch.dict(os.environ, clear=True)
+        self.env_mock.start()
+
+    def tearDown(self):
+        self.env_mock.stop()
+
     @mock.patch('opencensus.common.monitored_resource.monitored_resource'
                 '.gcp_metadata_config.GcpMetadataConfig')
-    def test_gcp_gce_monitored_resource(self, gcp_metadata_mock):
+    def test_gcp_gce_monitored_resource(self, gcp_md_mock):
         mocked_labels = {
             'instance_id': 'my-instance',
             'project_id': 'my-project',
             'zone': 'us-east1'
         }
 
-        gcp_metadata_mock.return_value = mock.Mock()
-        gcp_metadata_mock.return_value.get_gce_metadata.return_value =\
-            mocked_labels
+        gcp_md_mock.return_value = mock.Mock()
+        gcp_md_mock.return_value.get_gce_metadata.return_value = mocked_labels
+
         with mock_gce_env():
             resource = monitored_resource.get_instance()
         self.assertEqual(resource.get_type(), 'gce_instance')
-        self.assertDictEqual(resource.get_labels(), mocked_labels)
+        self.assertEqual(resource.get_labels(), mocked_labels)
+
+        with mock_oc_env():
+            with mock_gce_env():
+                resource = monitored_resource.get_instance()
+        self.assertEqual(resource.get_type(), 'mock_resource_type')
+        self.assertDictContainsSubset(
+            {'mock_label_key': 'mock_label_value'}, resource.get_labels())
+        self.assertDictContainsSubset(mocked_labels, resource.get_labels())
+
+    @mock.patch('opencensus.common.monitored_resource.monitored_resource'
+                '.gcp_metadata_config.GcpMetadataConfig')
+    def test_gcp_k8s_monitored_resource(self, gcp_md_mock):
 
-    def test_gcp_k8s_monitored_resource(self):
         mocked_labels = {
             'instance_id': 'my-instance',
             'cluster_name': 'cluster',
@@ -92,31 +119,53 @@ def test_gcp_k8s_monitored_resource(self):
             'namespace_id': 'namespace',
             'container_name': 'container'
         }
+        cluster_name_key = 'instance/attributes/cluster-name'
+        cluster_name_val = 'cluster'
+        gcp_md_mock.return_value = mock.Mock()
+        gcp_md_mock.return_value.get_gce_metadata.return_value = mocked_labels
+        gcp_md_mock.get_attribute.return_value = cluster_name_val
+
+        with mock_k8s_env():
+            r1 = monitored_resource.get_instance()
+
+        gcp_md_mock.get_attribute.assert_called_once_with(cluster_name_key)
+        self.assertEqual(r1.get_type(), 'k8s_container')
+        self.assertDictContainsSubset(mocked_labels, r1.get_labels())
 
-        with mock_mr_method('k8s_utils.get_k8s_metadata', mocked_labels):
+        with mock_oc_env():
             with mock_k8s_env():
-                resource = monitored_resource.get_instance()
-        self.assertEqual(resource.get_type(), 'k8s_container')
-        self.assertDictEqual(resource.get_labels(), mocked_labels)
+                r2 = monitored_resource.get_instance()
+
+        self.assertEqual(r1.get_type(), 'k8s_container')
+        self.assertDictContainsSubset(mocked_labels, r1.get_labels())
+        self.assertDictContainsSubset(
+            {'mock_label_key': 'mock_label_value'}, r2.get_labels())
 
     @mock.patch('opencensus.common.monitored_resource.monitored_resource'
                 '.aws_identity_doc_utils.AwsIdentityDocumentUtils')
-    def test_aws_monitored_resource(self, aws_metadata_mock):
+    def test_aws_monitored_resource(self, aws_md_mock):
 
         mocked_labels = {
             'instance_id': 'i-1234567890abcdef0',
             'aws_account': '123456789012',
             'region': 'us-west-2'
         }
 
-        aws_metadata_mock.return_value = mock.Mock()
-        aws_metadata_mock.return_value.get_aws_metadata.return_value =\
-            mocked_labels
+        aws_md_mock.return_value = mock.Mock()
+        aws_md_mock.return_value.get_aws_metadata.return_value = mocked_labels
 
         with mock_aws_env():
             resource = monitored_resource.get_instance()
         self.assertEqual(resource.get_type(), 'aws_ec2_instance')
-        self.assertDictEqual(resource.get_labels(), mocked_labels)
+        self.assertEqual(resource.get_labels(), mocked_labels)
+
+        with mock_oc_env():
+            with mock_aws_env():
+                resource = monitored_resource.get_instance()
+        self.assertEqual(resource.get_type(), 'mock_resource_type')
+        self.assertDictContainsSubset(
+            {'mock_label_key': 'mock_label_value'}, resource.get_labels())
+        self.assertDictContainsSubset(mocked_labels, resource.get_labels())
 
     def test_k8s_environment(self):
         patch = mock.patch.dict(os.environ,
@@ -147,7 +196,7 @@ def test_gce_environment(self):
                 'aws_identity_doc_utils.AwsIdentityDocumentUtils.'
                 'is_running_on_aws',
                 return_value=True)
-    def test_aws_environment(self, aws_util_mock, gcp_metadata_mock):
+    def test_aws_environment(self, aws_util_mock, gcp_md_mock):
         mr = monitored_resource.get_instance()
 
         self.assertIsNotNone(mr)
@@ -160,7 +209,13 @@ def test_aws_environment(self, aws_util_mock, gcp_metadata_mock):
                 'aws_identity_doc_utils.AwsIdentityDocumentUtils.'
                 'is_running_on_aws',
                 return_value=False)
-    def test_non_supported_environment(self, aws_util_mock, gcp_metadata_mock):
+    def test_non_supported_environment(self, aws_util_mock, gcp_md_mock):
         mr = monitored_resource.get_instance()
-
         self.assertIsNone(mr)
+
+        with mock_oc_env():
+            mr = monitored_resource.get_instance()
+        self.assertIsNotNone(mr)
+        self.assertEqual(mr.get_type(), 'mock_resource_type')
+        self.assertDictEqual(
+            mr.get_labels(), {'mock_label_key': 'mock_label_value'})

tests/unit/common/test_resource.py

@@ -19,6 +19,7 @@
 except ImportError:
     from unittest import mock
 
+import os
 import unittest
 
 from opencensus.common import resource as resource_module
@@ -123,6 +124,28 @@ def test_merge_overwrite(self):
 
 class TestResourceModule(unittest.TestCase):
 
+    def test_merge_resource(self):
+        with self.assertRaises(ValueError):
+            resource_module.merge_resources(None)
+        with self.assertRaises(ValueError):
+            resource_module.merge_resources([])
+
+        r1 = Resource(None, {'lk1': 'lv11'})
+        r2 = Resource('t2', {'lk1': 'lv12', 'lk2': 'lv22'})
+        r3 = Resource('t3', {'lk2': 'lv23', 'lk3': 'lv33'})
+
+        merged = resource_module.merge_resources([r1, r2, r3])
+        self.assertEqual(merged.type, 't2')
+        self.assertDictEqual(
+            merged.labels, {'lk1': 'lv11', 'lk2': 'lv22', 'lk3': 'lv33'})
+
+    def test_merge_resource_no_type(self):
+        r1 = Resource(None)
+        r2 = Resource(None)
+
+        merged = resource_module.merge_resources([r1, r2])
+        self.assertEqual(merged.type, None)
+
     def test_check_ascii_256(self):
         self.assertIsNone(resource_module.check_ascii_256(None))
 
@@ -156,12 +179,20 @@ def test_get_from_env_no_type(self):
         with mock.patch.dict('os.environ', {
                 'OC_RESOURCE_LABELS': 'k1=v1,k2=v2'
         }):
+            try:
+                del os.environ['OC_RESOURCE_TYPE']
+            except KeyError:
+                pass
             self.assertIsNone(resource_module.get_from_env())
 
     def test_get_from_env_no_labels(self):
         with mock.patch.dict('os.environ', {
                 'OC_RESOURCE_TYPE': 'opencensus.io/example',
         }):
+            try:
+                del os.environ['OC_RESOURCE_LABELS']
+            except KeyError:
+                pass
             resource = resource_module.get_from_env()
         self.assertEqual(resource.type, 'opencensus.io/example')
         self.assertDictEqual(resource.labels, {})