diff --git a/apps/ergw_core/src/ergw_pfcp_context.erl b/apps/ergw_core/src/ergw_pfcp_context.erl
index 8ceed7b4..b17693c7 100644
--- a/apps/ergw_core/src/ergw_pfcp_context.erl
+++ b/apps/ergw_core/src/ergw_pfcp_context.erl
@@ -990,6 +990,16 @@ select_upf_api_result(UPF, Session0, APNOpts) when is_binary(UPF) ->
 select_upf_api_result(_UPF, _Session, _APNOpts) ->
     {error, ?CTX_ERR(?FATAL, no_resources_available)}.
 
+session_set_pool(Pool, Value, Session) when is_binary(Value) ->
+    maps:put(Pool, Value, Session);
+session_set_pool(_, _, Session) ->
+    Session.
+
+session_set_maybe(true, Key, Value, Session) ->
+    maps:put(Key, Value, Session);
+session_set_maybe(_, _, _, Session) ->
+    Session.
+
 %% select_upf_3gpp/3
 select_upf_3gpp(Candidates, Session0, APNOpts) ->
     Wanted = maps:fold(fun apn_filter/3, #{}, APNOpts),
@@ -1002,12 +1012,11 @@ select_upf_3gpp(Candidates, Session0, APNOpts) ->
 	   PoolV6 = select_ip_pool(v6, IPvs, IPpools),
 	   NAT = select(random, NATblocks),
 
-	   Session1 = Session0#{'Framed-Pool' => PoolV4, 'Framed-IPv6-Pool' => PoolV6},
-	   Session2 = case maps:is_key(nat_port_blocks, Wanted) of
-			  true  -> Session1#{'NAT-Pool-Id' => NAT};
-			  false -> Session1
-		      end,
-	   Session = init_session_ue_ifid(APNOpts, Session2),
+	   Session1 = session_set_pool('Framed-Pool', PoolV4, Session0),
+	   Session2 = session_set_pool('Framed-IPv6-Pool', PoolV6, Session1),
+	   Session3 = session_set_maybe(maps:is_key(nat_port_blocks, Wanted),
+					 'NAT-Pool-Id', NAT, Session2),
+	   Session = init_session_ue_ifid(APNOpts, Session3),
 	   UPinfo = {Node, VRF, PoolV4, NAT, PoolV6},
 
 	   return({UPinfo, Session})