1 files changed, 23 insertions, 30 deletions

diff --git a/kamon-core-tests/src/test/scala/kamon/context/HttpPropagationSpec.scala b/kamon-core-tests/src/test/scala/kamon/context/HttpPropagationSpec.scala
index 44165b98..ac1250ea 100644
--- a/kamon-core-tests/src/test/scala/kamon/context/HttpPropagationSpec.scala
+++ b/kamon-core-tests/src/test/scala/kamon/context/HttpPropagationSpec.scala
@@ -2,7 +2,8 @@ package kamon.context
 
 import com.typesafe.config.ConfigFactory
 import kamon.Kamon
-import kamon.context.HttpPropagation.Direction
+import kamon.context.HttpPropagation.{HeaderReader, HeaderWriter}
+import kamon.context.Propagation.{EntryReader, EntryWriter}
 import org.scalatest.{Matchers, OptionValues, WordSpec}
 
 import scala.collection.mutable
@@ -12,9 +13,8 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
   "The HTTP Context Propagation" when {
     "reading from incoming requests" should {
       "return an empty context if there are no tags nor keys" in {
-        val context = httpPropagation.readContext(headerReaderFromMap(Map.empty))
-        context.tags shouldBe empty
-        context.entries shouldBe empty
+        val context = httpPropagation.read(headerReaderFromMap(Map.empty))
+        context.isEmpty() shouldBe true
       }
 
       "read tags from an HTTP message when they are available" in {
@@ -22,7 +22,7 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
           "x-content-tags" -> "hello=world;correlation=1234",
           "x-mapped-tag" -> "value"
         )
-        val context = httpPropagation.readContext(headerReaderFromMap(headers))
+        val context = httpPropagation.read(headerReaderFromMap(headers))
         context.tags should contain only(
           "hello" -> "world",
           "correlation" -> "1234",
@@ -32,7 +32,7 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
 
       "handle errors when reading HTTP headers" in {
         val headers = Map("fail" -> "")
-        val context = httpPropagation.readContext(headerReaderFromMap(headers))
+        val context = httpPropagation.read(headerReaderFromMap(headers))
         context.tags shouldBe empty
         context.entries shouldBe empty
       }
@@ -44,7 +44,7 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
           "integer-header" -> "123"
         )
 
-        val context = httpPropagation.readContext(headerReaderFromMap(headers))
+        val context = httpPropagation.read(headerReaderFromMap(headers))
         context.get(HttpPropagationSpec.StringKey) shouldBe "hey"
         context.get(HttpPropagationSpec.IntegerKey) shouldBe 123
         context.get(HttpPropagationSpec.OptionalKey) shouldBe empty
@@ -56,17 +56,9 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
 
 
     "writing to outgoing requests" should {
-      propagationWritingTests(Direction.Outgoing)
-    }
-
-    "writing to returning requests" should {
-      propagationWritingTests(Direction.Returning)
-    }
-
-    def propagationWritingTests(direction: Direction.Write) = {
       "not write anything if the context is empty" in {
         val headers = mutable.Map.empty[String, String]
-        httpPropagation.writeContext(Context.Empty, headerWriterFromMap(headers), direction)
+        httpPropagation.write(Context.Empty, headerWriterFromMap(headers))
         headers shouldBe empty
       }
 
@@ -77,7 +69,7 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
           "mappedTag" -> "value"
         ))
 
-        httpPropagation.writeContext(context, headerWriterFromMap(headers), direction)
+        httpPropagation.write(context, headerWriterFromMap(headers))
         headers should contain only(
           "x-content-tags" -> "hello=world;",
           "x-mapped-tag" -> "value"
@@ -91,10 +83,10 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
           HttpPropagationSpec.IntegerKey, 42,
         )
 
-        httpPropagation.writeContext(context, headerWriterFromMap(headers), direction)
+        httpPropagation.write(context, headerWriterFromMap(headers))
         headers should contain only(
           "string-header" -> "out-we-go"
-        )
+          )
       }
     }
   }
@@ -116,23 +108,24 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
         |entries.incoming.string = "kamon.context.HttpPropagationSpec$StringEntryCodec"
         |entries.incoming.integer = "kamon.context.HttpPropagationSpec$IntegerEntryCodec"
         |entries.outgoing.string = "kamon.context.HttpPropagationSpec$StringEntryCodec"
-        |entries.returning.string = "kamon.context.HttpPropagationSpec$StringEntryCodec"
         |
       """.stripMargin
     ).withFallback(ConfigFactory.load().getConfig("kamon.propagation")), Kamon)
 
 
   def headerReaderFromMap(map: Map[String, String]): HttpPropagation.HeaderReader = new HttpPropagation.HeaderReader {
-    override def readHeader(header: String): Option[String] = {
+    override def read(header: String): Option[String] = {
       if(map.get("fail").nonEmpty)
         sys.error("failing on purpose")
 
       map.get(header)
     }
+
+    override def readAll(): Map[String, String] = map
   }
 
   def headerWriterFromMap(map: mutable.Map[String, String]): HttpPropagation.HeaderWriter = new HttpPropagation.HeaderWriter {
-    override def writeHeader(header: String, value: String): Unit = map.put(header, value)
+    override def write(header: String, value: String): Unit = map.put(header, value)
   }
 }
 
@@ -143,23 +136,23 @@ object HttpPropagationSpec {
   val OptionalKey = Context.key[Option[String]]("optional", None)
 
 
-  class StringEntryCodec extends HttpPropagation.EntryReader with HttpPropagation.EntryWriter {
+  class StringEntryCodec extends EntryReader[HeaderReader] with EntryWriter[HeaderWriter] {
     private val HeaderName = "string-header"
 
-    override def readEntry(reader: HttpPropagation.HeaderReader, context: Context): Context = {
-      reader.readHeader(HeaderName)
+    override def read(reader: HttpPropagation.HeaderReader, context: Context): Context = {
+      reader.read(HeaderName)
         .map(v => context.withKey(StringKey, v))
         .getOrElse(context)
     }
 
-    override def writeEntry(context: Context, writer: HttpPropagation.HeaderWriter, direction: Direction.Write): Unit = {
-      Option(context.get(StringKey)).foreach(v => writer.writeHeader(HeaderName, v))
+    override def write(context: Context, writer: HttpPropagation.HeaderWriter): Unit = {
+      Option(context.get(StringKey)).foreach(v => writer.write(HeaderName, v))
     }
   }
 
-  class IntegerEntryCodec extends HttpPropagation.EntryReader {
-    override def readEntry(reader: HttpPropagation.HeaderReader, context: Context): Context = {
-      reader.readHeader("integer-header")
+  class IntegerEntryCodec extends EntryReader[HeaderReader] {
+    override def read(reader: HttpPropagation.HeaderReader, context: Context): Context = {
+      reader.read("integer-header")
         .map(v => context.withKey(IntegerKey, v.toInt))
         .getOrElse(context)