Skip to content

Commit

Permalink
Move DistributedContext implementation into api level.
Browse files Browse the repository at this point in the history
  • Loading branch information
a-feld committed Aug 26, 2019
1 parent 0c5e803 commit 55a0fda
Show file tree
Hide file tree
Showing 4 changed files with 42 additions and 62 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
# limitations under the License.

from contextlib import contextmanager
import abc
import typing

PRINTABLE = set(chr(num) for num in range(32, 127))
Expand Down Expand Up @@ -75,20 +74,23 @@ def __init__(
self.value = value


class DistributedContext(abc.ABC):
class DistributedContext(dict):
"""A container for distributed context entries"""

@abc.abstractmethod
def get_entries(self) -> typing.Iterable[Entry]:
"""Returns an immutable iterator to entries."""
return self.values()

@abc.abstractmethod
def get_entry_value(self, key: EntryKey) -> typing.Optional[EntryValue]:
def get_entry_value(
self,
key: EntryKey
) -> typing.Optional[EntryValue]:
"""Returns the entry associated with a key or None
Args:
key: the key with which to perform a lookup
"""
return self.get(key)


class DistributedContextManager:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,33 @@ def test_key_new(self):
self.assertEqual(key, "ok")


class TestDistributedContext(unittest.TestCase):
def setUp(self):
entry = self.entry = distributedcontext.Entry(
distributedcontext.EntryMetadata(
distributedcontext.EntryMetadata.NO_PROPAGATION,
),
distributedcontext.EntryKey("key"),
distributedcontext.EntryValue("value"),
)
context = self.context = distributedcontext.DistributedContext()
context[entry.key] = entry

def test_get_entries(self):
self.assertIn(self.entry, self.context.get_entries())

def test_get_entry_value_present(self):
value = self.context.get_entry_value(
self.entry.key,
)
self.assertIs(value, self.entry)

def test_get_entry_value_missing(self):
key = distributedcontext.EntryKey("missing")
value = self.context.get_entry_value(key)
self.assertIsNone(value)


class TestDistributedContextManager(unittest.TestCase):
def setUp(self):
self.manager = distributedcontext.DistributedContextManager()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,25 +19,6 @@
from opentelemetry.context import Context


class DistributedContext(dict, dctx_api.DistributedContext):
"""A container for distributed context entries"""

def get_entries(self) -> typing.Iterable[dctx_api.Entry]:
"""Returns an immutable iterator to entries."""
return self.values()

def get_entry_value(
self,
key: dctx_api.EntryKey
) -> typing.Optional[dctx_api.EntryValue]:
"""Returns the entry associated with a key or None
Args:
key: the key with which to perform a lookup
"""
return self.get(key)


class DistributedContextManager(dctx_api.DistributedContextManager):
"""See `opentelemetry.distributedcontext.DistributedContextManager`
Expand All @@ -53,7 +34,9 @@ def __init__(self, name: str = "") -> None:

self._current_context = Context.register_slot(slot_name)

def get_current_context(self) -> typing.Optional[DistributedContext]:
def get_current_context(
self,
) -> typing.Optional[dctx_api.DistributedContext]:
"""Gets the current DistributedContext.
Returns:
Expand All @@ -64,8 +47,8 @@ def get_current_context(self) -> typing.Optional[DistributedContext]:
@contextmanager
def use_context(
self,
context: DistributedContext,
) -> typing.Iterator[DistributedContext]:
context: dctx_api.DistributedContext,
) -> typing.Iterator[dctx_api.DistributedContext]:
"""Context manager for controlling a DistributedContext lifetime.
Set the context as the active DistributedContext.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,42 +14,10 @@

import unittest


from opentelemetry.distributedcontext import (
Entry,
EntryMetadata,
EntryKey,
EntryValue,
)

from opentelemetry import distributedcontext as dctx_api
from opentelemetry.sdk import distributedcontext


class TestDistributedContext(unittest.TestCase):
def setUp(self):
entry = self.entry = Entry(
EntryMetadata(EntryMetadata.NO_PROPAGATION),
EntryKey("key"),
EntryValue("value"),
)
context = self.context = distributedcontext.DistributedContext()
context[entry.key] = entry

def test_get_entries(self):
self.assertIn(self.entry, self.context.get_entries())

def test_get_entry_value_present(self):
value = self.context.get_entry_value(
self.entry.key,
)
self.assertIs(value, self.entry)

def test_get_entry_value_missing(self):
key = EntryKey("missing")
value = self.context.get_entry_value(key)
self.assertIsNone(value)


class TestDistributedContextManager(unittest.TestCase):
def setUp(self):
self.manager = distributedcontext.DistributedContextManager()
Expand All @@ -59,7 +27,7 @@ def test_use_context(self):
self.assertIsNone(self.manager.get_current_context())

# Start initial context
dctx = distributedcontext.DistributedContext()
dctx = dctx_api.DistributedContext()
with self.manager.use_context(dctx) as current:
self.assertIs(current, dctx)
self.assertIs(
Expand All @@ -68,7 +36,7 @@ def test_use_context(self):
)

# Context is overridden
nested_dctx = distributedcontext.DistributedContext()
nested_dctx = dctx_api.DistributedContext()
with self.manager.use_context(nested_dctx) as current:
self.assertIs(current, nested_dctx)
self.assertIs(
Expand Down

0 comments on commit 55a0fda

Please sign in to comment.