diff --git a/proto/marshal/collections/maps.py b/proto/marshal/collections/maps.py
index a66e8fc5..3c485716 100644
--- a/proto/marshal/collections/maps.py
+++ b/proto/marshal/collections/maps.py
@@ -65,15 +65,8 @@ def __setitem__(self, key, value):
         #
         # Therefore, we create a key that way (clearing any fields that may
         # be set) and then merge in our values.
-        # TODO: self.pb[key] should always be Message. Remove this after protobuf fix.
-        # In UPB, sometimes self.pb[key] is not always a proto.
-        # This happens during marshalling when the pb_value is upb.MapCompositeContainer
-        # so it's not marshalled correcrtly (i.e. should be scalar values not composite).
-        if isinstance(self.pb[key], Message):
-            self.pb[key].Clear()
-            self.pb[key].MergeFrom(pb_value)
-        else:
-            self.pb[key] = value
+        self.pb[key].Clear()
+        self.pb[key].MergeFrom(pb_value)
 
     def __delitem__(self, key):
         self.pb.pop(key)
diff --git a/tests/test_marshal_types_enum.py b/tests/test_marshal_types_enum.py
index 02626a8d..acb55461 100644
--- a/tests/test_marshal_types_enum.py
+++ b/tests/test_marshal_types_enum.py
@@ -97,5 +97,4 @@ class MolluscContainer(proto.Message):
     clam = Bivalve.CLAM
     mc.bivalves["clam"] = clam
     mc.bivalves["oyster"] = 1
-
-    assert mc.bivalves == {"clam": clam, "oyster": Bivalve.OYSTER}
+    assert dict(mc.bivalves) == {"clam": clam, "oyster": Bivalve.OYSTER}