1313# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1414# See the License for the specific language governing permissions and
1515# limitations under the License.
16+ from unittest .mock import Mock
1617
1718from math import inf , nan
1819
2122 encode_pretty_printed_json ,
2223 iterencode_canonical_json ,
2324 iterencode_pretty_printed_json ,
24- set_json_library ,
25+ set_json_library , register_preserialisation_callback ,
2526)
2627
2728import unittest
@@ -156,4 +157,32 @@ class C:
156157 pass
157158
158159 with self .assertRaises (Exception ):
159- encode_canonical_json (C ())
160+ encode_canonical_json (C ())
161+
162+ def test_preserialisation_callback (self ) -> None :
163+ class C :
164+ pass
165+
166+ register_preserialisation_callback (C , lambda c : "I am a C instance" )
167+
168+ result = encode_canonical_json (C ())
169+ self .assertEqual (result , b'"I am a C instance"' )
170+
171+ def test_cannot_register_preserialisation_callback_for_object (self ) -> None :
172+ with self .assertRaises (Exception ):
173+ register_preserialisation_callback (object , lambda c : "shouldn't be able to do this" )
174+
175+ def test_most_recent_preserialisation_callback_called (self ) -> None :
176+ class C :
177+ pass
178+
179+ callback1 = Mock (return_value = "callback 1 was called" )
180+ callback2 = Mock (return_value = "callback 2 was called" )
181+
182+ register_preserialisation_callback (C , callback1 )
183+ register_preserialisation_callback (C , callback2 )
184+
185+ encode_canonical_json (C ())
186+
187+ callback1 .assert_not_called ()
188+ callback2 .assert_called_once ()
0 commit comments