diff --git a/src/contract_driver.erl b/src/contract_driver.erl
index 0e87569..8953838 100644
--- a/src/contract_driver.erl
+++ b/src/contract_driver.erl
@@ -33,8 +33,8 @@
 -type contract() :: module().
 -type options() :: list(term()).
 -type parsed_options() :: term().
--type cont_init() :: {init,Rest::term(),Extras::term()} | {more,More::fun()}.
--type cont_done() :: {done,Term::term(),Rest::term(),Extras::term()} | {more,More::fun()}.
+-type cont_init() :: {init,Rest::term(),Extras::term()} | {more,More::term()}.
+-type cont_done() :: {done,Term::term(),Rest::term(),Extras::term()} | {more,More::term()}.
 -type io() :: any().
 
 -callback start(contract()) -> pid().
diff --git a/test/qc/qc_ubf_gen_tests.erl b/test/qc/qc_ubf_gen_tests.erl
index 16e9803..84be62e 100644
--- a/test/qc/qc_ubf_gen_tests.erl
+++ b/test/qc/qc_ubf_gen_tests.erl
@@ -64,7 +64,7 @@ prop_ubf_gen_any() ->
     ?FORALL(X, qc_gen:qc_any(),
             begin
                 UBF = ubf:encode(X),
-                {ok, Y, ""} = ubf:decode(UBF),
+                {done, Y, "", undefined} = ubf:decode(UBF),
                 Res = X =:= Y,
                 ?WHENFAIL(io:format("~n~p:~p ~p -> ~p -> ~p~n",[?FILE, ?LINE, X, UBF, Y]), Res)
             end).