Check for null in TagContextBuilderImpl.remove and add tests.

2 files changed, 29 insertions, 1 deletions

diff --git a/impl_core/src/main/java/io/opencensus/implcore/tags/TagContextBuilderImpl.java b/impl_core/src/main/java/io/opencensus/implcore/tags/TagContextBuilderImpl.java
index 29ff9a3b..a17198d8 100644
--- a/impl_core/src/main/java/io/opencensus/implcore/tags/TagContextBuilderImpl.java
+++ b/impl_core/src/main/java/io/opencensus/implcore/tags/TagContextBuilderImpl.java
@@ -44,7 +44,7 @@ final class TagContextBuilderImpl extends TagContextBuilder {
 
   @Override
   public TagContextBuilderImpl remove(TagKey key) {
-    tags.remove(key);
+    tags.remove(checkNotNull(key, "key"));
     return this;
   }
 
diff --git a/impl_core/src/test/java/io/opencensus/implcore/tags/TagContextImplTest.java b/impl_core/src/test/java/io/opencensus/implcore/tags/TagContextImplTest.java
index 1dd43da2..2ad4ea73 100644
--- a/impl_core/src/test/java/io/opencensus/implcore/tags/TagContextImplTest.java
+++ b/impl_core/src/test/java/io/opencensus/implcore/tags/TagContextImplTest.java
@@ -25,6 +25,7 @@ import com.google.common.collect.Lists;
 import com.google.common.testing.EqualsTester;
 import io.opencensus.tags.Tag;
 import io.opencensus.tags.TagContext;
+import io.opencensus.tags.TagContextBuilder;
 import io.opencensus.tags.TagKey;
 import io.opencensus.tags.TagValue;
 import io.opencensus.tags.Tagger;
@@ -81,6 +82,24 @@ public class TagContextImplTest {
   }
 
   @Test
+  public void put_nullKey() {
+    TagContext tags = new TagContextImpl(ImmutableMap.of(K1, V1));
+    TagContextBuilder builder = tagger.toBuilder(tags);
+    thrown.expect(NullPointerException.class);
+    thrown.expectMessage("key");
+    builder.put(null, V2);
+  }
+
+  @Test
+  public void put_nullValue() {
+    TagContext tags = new TagContextImpl(ImmutableMap.of(K1, V1));
+    TagContextBuilder builder = tagger.toBuilder(tags);
+    thrown.expect(NullPointerException.class);
+    thrown.expectMessage("value");
+    builder.put(K2, null);
+  }
+
+  @Test
   public void remove_existingKey() {
     TagContext tags = new TagContextImpl(ImmutableMap.of(K1, V1, K2, V2));
     assertThat(((TagContextImpl) tagger.toBuilder(tags).remove(K1).build()).getTags())
@@ -95,6 +114,15 @@ public class TagContextImplTest {
   }
 
   @Test
+  public void remove_nullValue() {
+    TagContext tags = new TagContextImpl(ImmutableMap.of(K1, V1));
+    TagContextBuilder builder = tagger.toBuilder(tags);
+    thrown.expect(NullPointerException.class);
+    thrown.expectMessage("key");
+    builder.remove(null);
+  }
+
+  @Test
   public void testIterator() {
     TagContextImpl tags = new TagContextImpl(ImmutableMap.of(K1, V1, K2, V2));
     Iterator<Tag> i = tags.getIterator();