From d1130795068a7f265f4d409b2f7190bd42b7b478 Mon Sep 17 00:00:00 2001
From: Diego Hurtado <ocelotl@users.noreply.github.com>
Date: Mon, 26 Apr 2021 16:39:42 -0600
Subject: [PATCH] Use traceflags object

Fixes #471
---
 .../src/opentelemetry/propagators/ot_trace/__init__.py    | 2 +-
 .../tests/test_ot_trace_propagator.py                     | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/propagator/opentelemetry-propagator-ot-trace/src/opentelemetry/propagators/ot_trace/__init__.py b/propagator/opentelemetry-propagator-ot-trace/src/opentelemetry/propagators/ot_trace/__init__.py
index 4013237749..0097a1bcfb 100644
--- a/propagator/opentelemetry-propagator-ot-trace/src/opentelemetry/propagators/ot_trace/__init__.py
+++ b/propagator/opentelemetry-propagator-ot-trace/src/opentelemetry/propagators/ot_trace/__init__.py
@@ -85,7 +85,7 @@ def extract(
                         trace_id=int(traceid, 16),
                         span_id=int(spanid, 16),
                         is_remote=True,
-                        trace_flags=traceflags,
+                        trace_flags=TraceFlags(traceflags),
                     )
                 ),
                 context,
diff --git a/propagator/opentelemetry-propagator-ot-trace/tests/test_ot_trace_propagator.py b/propagator/opentelemetry-propagator-ot-trace/tests/test_ot_trace_propagator.py
index 5afa91d5da..f8b0a2b60a 100644
--- a/propagator/opentelemetry-propagator-ot-trace/tests/test_ot_trace_propagator.py
+++ b/propagator/opentelemetry-propagator-ot-trace/tests/test_ot_trace_propagator.py
@@ -248,6 +248,10 @@ def test_extract_trace_id_span_id_sampled_true(self):
         self.assertEqual(hex(span_context.span_id)[2:], "e457b5a2e4d86bd1")
         self.assertTrue(span_context.is_remote)
         self.assertEqual(span_context.trace_flags, TraceFlags.SAMPLED)
+        self.assertIsInstance(
+            get_current_span().get_span_context().trace_flags,
+            TraceFlags
+        )
 
     def test_extract_trace_id_span_id_sampled_false(self):
         """Test valid trace_id, span_id and sampled false"""
@@ -268,6 +272,10 @@ def test_extract_trace_id_span_id_sampled_false(self):
         self.assertEqual(hex(span_context.span_id)[2:], "e457b5a2e4d86bd1")
         self.assertTrue(span_context.is_remote)
         self.assertEqual(span_context.trace_flags, TraceFlags.DEFAULT)
+        self.assertIsInstance(
+            get_current_span().get_span_context().trace_flags,
+            TraceFlags
+        )
 
     def test_extract_malformed_trace_id(self):
         """Test extraction with malformed trace_id"""