Fix codec

wvlet · Jan 21, 2024 · f6d6d15 · f6d6d15
1 parent f83e4e4
commit f6d6d15
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 3 deletions.
diff --git a/airframe-codec/src/main/scala/wvlet/airframe/codec/MessageCodecFactory.scala b/airframe-codec/src/main/scala/wvlet/airframe/codec/MessageCodecFactory.scala
@@ -11,7 +11,10 @@ case class MessageCodecFactory(codecFinder: MessageCodecFinder = Compat.messageC
   private[this] var cache = Map.empty[String, MessageCodec[_]]
 
   def withCodecs(additionalCodecs: Map[Surface, MessageCodec[_]]): MessageCodecFactory = {
-    this.copy(codecFinder = codecFinder orElse MessageCodecFinder.newCodecFinder(additionalCodecs))
+    this.copy(codecFinder = MessageCodecFinder.newCodecFinder(additionalCodecs) orElse codecFinder)
+  }
+  def withCodecs(additionalCodecs: PartialFunction[Surface, MessageCodec[_]]): MessageCodecFactory = {
+    this.copy(codecFinder = MessageCodecFinder.newCodecFinder(additionalCodecs) orElse codecFinder)
   }
 
   // Generate a codec that outputs objects as Map type. This should be enabled for generating JSON data

diff --git a/airframe-codec/src/main/scala/wvlet/airframe/codec/MessageCodecFinder.scala b/airframe-codec/src/main/scala/wvlet/airframe/codec/MessageCodecFinder.scala
@@ -56,6 +56,14 @@ object MessageCodecFinder extends LogSupport {
       }
     }
 
+  def newCodecFinder(codecTable: PartialFunction[Surface, MessageCodec[_]]): MessageCodecFinder =
+    new MessageCodecFinder {
+      override def findCodec(
+          factory: MessageCodecFactory,
+          seenSet: Set[Surface]
+      ): PartialFunction[Surface, MessageCodec[_]] = codecTable
+    }
+
   val defaultKnownCodecs: Map[Surface, MessageCodec[_]] = {
     StandardCodec.standardCodec ++
       MetricsCodec.metricsCodec ++

diff --git a/airframe-http-netty/src/main/scala/wvlet/airframe/http/netty/NettyServer.scala b/airframe-http-netty/src/main/scala/wvlet/airframe/http/netty/NettyServer.scala
@@ -238,7 +238,7 @@ class NettyServer(config: NettyServerConfig, session: Session) extends HttpServe
               NettyBackend,
               new NettyResponseHandler,
               // Set a custom codec and use JSON map output
-              MessageCodecFactory.newFactory(config.customCodec).withMapOutput
+              MessageCodecFactory.defaultFactoryForJSON.withCodecs(config.customCodec)
             )
           )
       }

diff --git a/airframe-http-netty/src/test/scala/wvlet/airframe/http/netty/NettyRESTServerTest.scala b/airframe-http-netty/src/test/scala/wvlet/airframe/http/netty/NettyRESTServerTest.scala
@@ -157,7 +157,7 @@ class NettyRESTServerTest extends AirSpec {
         client.send(Http.GET("/v1/rich_info")).map { response => response.contentString }
       }
       Rx.fromSeq(futures).map { result =>
-        info(result.mkString(", "))
+        info(result)
       }
     }
 //