2 files changed, 19 insertions, 14 deletions
diff --git a/kamon-core-tests/src/test/scala/kamon/context/BinaryPropagationSpec.scala b/kamon-core-tests/src/test/scala/kamon/context/BinaryPropagationSpec.scala
index 5681d300..4fa7116d 100644
--- a/kamon-core-tests/src/test/scala/kamon/context/BinaryPropagationSpec.scala
+++ b/kamon-core-tests/src/test/scala/kamon/context/BinaryPropagationSpec.scala
@@ -6,6 +6,8 @@ import com.typesafe.config.ConfigFactory
 import kamon.Kamon
 import kamon.context.BinaryPropagation.{ByteStreamReader, ByteStreamWriter}
 import kamon.context.Propagation.{EntryReader, EntryWriter}
+import kamon.tag.TagSet
+import kamon.tag.Lookups._
 import org.scalatest.{Matchers, OptionValues, WordSpec}
 
 import scala.util.Random
@@ -70,13 +72,14 @@ class BinaryPropagationSpec extends WordSpec with Matchers with OptionValues {
     }
 
     "round trip a Context that only has tags" in {
-      val context = Context.of(Map("hello" -> "world", "kamon" -> "rulez"))
+      val context = Context.of(TagSet.from(Map("hello" -> "world", "kamon" -> "rulez")))
       val writer = inspectableByteStreamWriter()
       binaryPropagation.write(context, writer)
 
       val rtContext = binaryPropagation.read(ByteStreamReader.of(writer.toByteArray))
       rtContext.entries shouldBe empty
-      rtContext.tags should contain theSameElementsAs (context.tags)
+      rtContext.tags.get(plain("hello")) shouldBe "world"
+      rtContext.tags.get(plain("kamon")) shouldBe "rulez"
     }
 
     "round trip a Context that only has entries" in {
@@ -91,7 +94,7 @@ class BinaryPropagationSpec extends WordSpec with Matchers with OptionValues {
     }
 
     "round trip a Context that with tags and entries" in {
-      val context = Context.of(Map("hello" -> "world", "kamon" -> "rulez"))
+      val context = Context.of(TagSet.from(Map("hello" -> "world", "kamon" -> "rulez")))
         .withKey(BinaryPropagationSpec.StringKey, "string-value")
         .withKey(BinaryPropagationSpec.IntegerKey, 42)
 
@@ -99,7 +102,8 @@ class BinaryPropagationSpec extends WordSpec with Matchers with OptionValues {
       binaryPropagation.write(context, writer)
       val rtContext = binaryPropagation.read(ByteStreamReader.of(writer.toByteArray))
 
-      rtContext.tags should contain theSameElementsAs (context.tags)
+      rtContext.tags.get(plain("hello")) shouldBe "world"
+      rtContext.tags.get(plain("kamon")) shouldBe "rulez"
       rtContext.get(BinaryPropagationSpec.StringKey) shouldBe "string-value"
       rtContext.get(BinaryPropagationSpec.IntegerKey) shouldBe 0 // there is no entry configuration for the integer key
     }
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 fcddfe24..0cd10672 100644
--- a/kamon-core-tests/src/test/scala/kamon/context/HttpPropagationSpec.scala
+++ b/kamon-core-tests/src/test/scala/kamon/context/HttpPropagationSpec.scala
@@ -5,6 +5,8 @@ import kamon.Kamon
 import kamon.context.HttpPropagation.{HeaderReader, HeaderWriter}
 import kamon.context.Propagation.{EntryReader, EntryWriter}
 import org.scalatest.{Matchers, OptionValues, WordSpec}
+import kamon.tag.Lookups._
+import kamon.tag.TagSet
 
 import scala.collection.mutable
 
@@ -22,12 +24,11 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
           "x-content-tags" -> "hello=world;correlation=1234",
           "x-mapped-tag" -> "value"
         )
+
         val context = httpPropagation.read(headerReaderFromMap(headers))
-        context.tags should contain only(
-          "hello" -> "world",
-          "correlation" -> "1234",
-          "mappedTag" -> "value"
-        )
+        context.tags.get(plain("hello")) shouldBe "world"
+        context.tags.get(plain("correlation")) shouldBe "1234"
+        context.tags.get(plain("mappedTag")) shouldBe "value"
       }
 
       "handle errors when reading HTTP headers" in {
@@ -48,9 +49,9 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
         context.get(HttpPropagationSpec.StringKey) shouldBe "hey"
         context.get(HttpPropagationSpec.IntegerKey) shouldBe 123
         context.get(HttpPropagationSpec.OptionalKey) shouldBe empty
-        context.getTag("hello").value shouldBe "world"
-        context.getTag("correlation").value shouldBe "1234"
-        context.getTag("unknown") shouldBe empty
+        context.getTag(plain("hello")) shouldBe "world"
+        context.getTag(option("correlation")).value shouldBe "1234"
+        context.getTag(option("unknown")) shouldBe empty
       }
     }
 
@@ -64,10 +65,10 @@ class HttpPropagationSpec extends WordSpec with Matchers with OptionValues {
 
       "write context tags when available" in {
         val headers = mutable.Map.empty[String, String]
-        val context = Context.of(Map(
+        val context = Context.of(TagSet.from(Map(
           "hello" -> "world",
           "mappedTag" -> "value"
-        ))
+        )))
 
         httpPropagation.write(context, headerWriterFromMap(headers))
         headers should contain only(