[tf.data] Adds more coverage for Defun-defun switch, adds map unit test.