fixing documentation issues and python2 testing issue

tswast · Aug 24, 2020 · Aug 1, 2020 · Aug 3, 2020 · Aug 4, 2020 · Aug 4, 2020
commit 2bcd106876db8b552b274b94d9497ea5c98f8cd7
diff --git a/google/cloud/bigquery/opentelemetry_tracing.py b/google/cloud/bigquery/opentelemetry_tracing.py
@@ -34,32 +34,13 @@
 
     HAS_OPENTELEMETRY = False
 
-
 _default_attributes = {
     "db.system": "BigQuery",
 }
 
 
 @contextmanager
 def create_span(name, attributes=None, client=None, job_ref=None):
-    # yield new span value
-    if not HAS_OPENTELEMETRY:
-        yield None
-        return
-    tracer = trace.get_tracer(__name__)
-
-    final_attributes = _get_final_span_attributes(attributes, client, job_ref)
-    with tracer.start_as_current_span(name=name, attributes=final_attributes) as span:
-        try:
-            yield span
-            span.set_status(Status(http_status_to_canonical_code(200)))
-        except GoogleAPICallError as error:
-            if error.code is not None:
-                span.set_status(Status(http_status_to_canonical_code(error.code)))
-            raise
-
-
-def _get_final_span_attributes(attributes=None, client=None, job_ref=None):
     """Creates a ContextManager for a Span to be exported to the configured exporter. If no configuration
         exists yields None.
 
@@ -83,7 +64,24 @@ def _get_final_span_attributes(attributes=None, client=None, job_ref=None):
                     Raised if a span could not be yielded or issue with call to
                     OpenTelemetry.
             """
+    final_attributes = _get_final_span_attributes(attributes, client, job_ref)
+    if not HAS_OPENTELEMETRY:
+        yield None
+        return
+    tracer = trace.get_tracer(__name__)
+
+    # yield new span value
+    with tracer.start_as_current_span(name=name, attributes=final_attributes) as span:
+        try:
+            yield span
+            span.set_status(Status(http_status_to_canonical_code(200)))
+        except GoogleAPICallError as error:
+            if error.code is not None:
+                span.set_status(Status(http_status_to_canonical_code(error.code)))
+            raise
 
+
+def _get_final_span_attributes(attributes=None, client=None, job_ref=None):
     if client:
         client_attributes = _set_client_attributes(client)
         _default_attributes.update(client_attributes)

diff --git a/noxfile.py b/noxfile.py
@@ -33,7 +33,7 @@ def default(session):
     """
     # Install all test dependencies, then install local packages in-place.
     session.install(
-        "mock", "pytest", "google-cloud-testutils", "pytest-cov", "freezegun",
+        "mock", "pytest", "google-cloud-testutils", "pytest-cov", "freezegun"
     )
     session.install("grpcio")
 

diff --git a/tests/unit/test_opentelemetry_tracing.py b/tests/unit/test_opentelemetry_tracing.py
@@ -1,20 +1,26 @@
 import sys
 import pytest
 
-from importlib import reload
 import mock
 
 from google.cloud.bigquery import opentelemetry_tracing, client
 
-from opentelemetry import trace
-from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor
-from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+try:
+    import opentelemetry
+    from opentelemetry import trace
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor
+    from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+except ImportError:
+    opentelemetry = None
+
+from six.moves import reload_module
 
 TEST_SPAN_NAME = "bar"
 TEST_SPAN_ATTRIBUTES = {"foo": "baz"}
 
 
+@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`")
 @pytest.fixture
 def setup():
     tracer_provider = TracerProvider()
@@ -25,65 +31,51 @@ def setup():
     yield memory_exporter
 
 
+@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`")
 def test_opentelemetry_not_installed(setup):
     temp_module = sys.modules["opentelemetry"]
     sys.modules["opentelemetry"] = None
-    reload(opentelemetry_tracing)
+    reload_module(opentelemetry_tracing)
     with opentelemetry_tracing.create_span("No-op for opentelemetry") as span:
         assert span is None
     sys.modules["opentelemetry"] = temp_module
-    reload(opentelemetry_tracing)
+    reload_module(opentelemetry_tracing)
 
 
+@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`")
 def test_opentelemetry_success(setup):
     expected_attributes = {"foo": "baz", "db.system": "BigQuery"}
+
     with opentelemetry_tracing.create_span(
-        TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES,
+            TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=None, job_ref=None
     ) as span:
-        if span is None:
-            span_list = setup.get_finished_spans()
-            assert len(span_list) == 1
+        assert span is not None
         assert span.name == TEST_SPAN_NAME
         assert span.attributes == expected_attributes
 
 
+@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`")
 def test_default_client_attributes(setup):
-    import google.auth.credentials
-
-    mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials)
-    test_client = client.Client(
-        project="test_project", credentials=mock_credentials, location="test_location"
-    )
 
     expected_attributes = {
         "foo": "baz",
         "db.system": "BigQuery",
         "db.name": "test_project",
         "location": "test_location",
     }
-    with opentelemetry_tracing.create_span(
-        TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client
-    ) as span:
-        if span is None:
-            span_list = setup.get_finished_spans()
-            assert len(span_list) == 1
-        assert span.name == TEST_SPAN_NAME
-        assert span.attributes == expected_attributes
-
-
+    with mock.patch('google.cloud.bigquery.client.Client') as test_client:
+        test_client.project = 'test_project'
+        test_client.location = 'test_location'
+        with opentelemetry_tracing.create_span(
+                TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, client=test_client
+        ) as span:
+            assert span is not None
+            assert span.name == TEST_SPAN_NAME
+            assert span.attributes == expected_attributes
+
+
+@pytest.mark.skipif(opentelemetry is None, reason="Require `opentelemetry`")
 def test_default_job_attributes(setup):
-    from google.cloud.bigquery import job
-    import google.auth.credentials
-
-    mock_credentials = mock.Mock(spec=google.auth.credentials.Credentials)
-
-    test_job_reference = job._JobReference(
-        job_id="test_job_id", project="test_project_id", location="test_location"
-    )
-    test_client = client.Client(
-        project="test_project", credentials=mock_credentials, location="test_location"
-    )
-    test_job = job._AsyncJob(job_id=test_job_reference, client=test_client)
 
     expected_attributes = {
         "db.system": "BigQuery",
@@ -93,12 +85,17 @@ def test_default_job_attributes(setup):
         "job_id": "test_job_id",
         "foo": "baz",
     }
+    with mock.patch('google.cloud.bigquery.job._AsyncJob') as test_job_ref:
+        test_job_ref.job_id = 'test_job_id'
+        test_job_ref.location = 'test_location'
+        test_job_ref.project = 'test_project_id'
+        test_job_ref.num_child_jobs = '0'
+
+        with opentelemetry_tracing.create_span(
+                TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job_ref
+        ) as span:
+            assert span is not None
+            assert span.name == TEST_SPAN_NAME
+            assert span.attributes == expected_attributes
+
 
-    with opentelemetry_tracing.create_span(
-        TEST_SPAN_NAME, attributes=TEST_SPAN_ATTRIBUTES, job_ref=test_job
-    ) as span:
-        if span is None:
-            span_list = setup.get_finished_spans()
-            assert len(span_list) == 1
-        assert span.name == TEST_SPAN_NAME
-        assert span.attributes == expected_attributes