diff --git a/cloud/agent/e2e/main_test.go b/cloud/agent/e2e/main_test.go
index 53b719b..9dfd129 100644
--- a/cloud/agent/e2e/main_test.go
+++ b/cloud/agent/e2e/main_test.go
@@ -115,7 +115,7 @@
 				Host:       registryAddr.String(),
 				Repository: "testos",
 				Tag:        "latest",
-				Digest:     image.ManifestDigest,
+				Digest:     image.Digest(),
 			},
 			NodeParameters: &mpb.NodeParameters{},
 			RootDevice:     "vda",
@@ -179,7 +179,7 @@
 	grpcListenAddr := grpcLis.Addr().(*net.TCPAddr)
 
 	registryServer := registry.NewServer()
-	registryServer.AddImage("testos", "latest", image)
+	registryServer.AddRef("testos", "latest", image)
 
 	registryLis, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {