diff --git a/src/kdf_test.go b/src/kdf_test.go
index 0cce81d..a89dacc 100644
--- a/src/kdf_test.go
+++ b/src/kdf_test.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"encoding/hex"
+	"golang.org/x/crypto/blake2s"
 	"testing"
 )
 
@@ -44,10 +45,12 @@ func TestKDF(t *testing.T) {
 		},
 	}
 
+	var t0, t1, t2 [blake2s.Size]byte
+
 	for _, test := range tests {
 		key, _ := hex.DecodeString(test.key)
 		input, _ := hex.DecodeString(test.input)
-		t0, t1, t2 := KDF3(key, input)
+		KDF3(&t0, &t1, &t2, key, input)
 		t0s := hex.EncodeToString(t0[:])
 		t1s := hex.EncodeToString(t1[:])
 		t2s := hex.EncodeToString(t2[:])
@@ -59,7 +62,7 @@ func TestKDF(t *testing.T) {
 	for _, test := range tests {
 		key, _ := hex.DecodeString(test.key)
 		input, _ := hex.DecodeString(test.input)
-		t0, t1 := KDF2(key, input)
+		KDF2(&t0, &t1, key, input)
 		t0s := hex.EncodeToString(t0[:])
 		t1s := hex.EncodeToString(t1[:])
 		assertEquals(t, t0s, test.t0)
@@ -69,7 +72,7 @@ func TestKDF(t *testing.T) {
 	for _, test := range tests {
 		key, _ := hex.DecodeString(test.key)
 		input, _ := hex.DecodeString(test.input)
-		t0 := KDF1(key, input)
+		KDF1(&t0, key, input)
 		t0s := hex.EncodeToString(t0[:])
 		assertEquals(t, t0s, test.t0)
 	}
diff --git a/src/noise_test.go b/src/noise_test.go
index 0d7f0e9..48408f9 100644
--- a/src/noise_test.go
+++ b/src/noise_test.go
@@ -117,8 +117,8 @@ func TestNoiseHandshake(t *testing.T) {
 		var err error
 		var out []byte
 		var nonce [12]byte
-		out = key1.send.Seal(out, nonce[:], testMsg, nil)
-		out, err = key2.receive.Open(out[:0], nonce[:], out, nil)
+		out = key1.send.aead.Seal(out, nonce[:], testMsg, nil)
+		out, err = key2.receive.aead.Open(out[:0], nonce[:], out, nil)
 		assertNil(t, err)
 		assertEqual(t, out, testMsg)
 	}()
@@ -128,8 +128,8 @@ func TestNoiseHandshake(t *testing.T) {
 		var err error
 		var out []byte
 		var nonce [12]byte
-		out = key2.send.Seal(out, nonce[:], testMsg, nil)
-		out, err = key1.receive.Open(out[:0], nonce[:], out, nil)
+		out = key2.send.aead.Seal(out, nonce[:], testMsg, nil)
+		out, err = key1.receive.aead.Open(out[:0], nonce[:], out, nil)
 		assertNil(t, err)
 		assertEqual(t, out, testMsg)
 	}()