diff --git a/pkg/plugin/implementation/keymanager/cmd/plugin.go b/pkg/plugin/implementation/keymanager/cmd/plugin.go
index 5e37af4..b4450d3 100644
--- a/pkg/plugin/implementation/keymanager/cmd/plugin.go
+++ b/pkg/plugin/implementation/keymanager/cmd/plugin.go
@@ -9,9 +9,7 @@ import (
 )
 
 // keyManagerProvider implements the plugin provider for the KeyManager plugin.
-type keyManagerProvider struct {
-	newFunc func(ctx context.Context, cache definition.Cache, registry definition.RegistryLookup, cfg *keymanager.Config) (definition.KeyManager, func() error, error)
-}
+type keyManagerProvider struct{}
 
 // newKeyManagerFunc is a function type that creates a new KeyManager instance.
 var newKeyManagerFunc = keymanager.New
diff --git a/pkg/plugin/implementation/keymanager/cmd/plugin_test.go b/pkg/plugin/implementation/keymanager/cmd/plugin_test.go
index bec7c6f..881b837 100644
--- a/pkg/plugin/implementation/keymanager/cmd/plugin_test.go
+++ b/pkg/plugin/implementation/keymanager/cmd/plugin_test.go
@@ -11,43 +11,6 @@ import (
 	"github.com/beckn/beckn-onix/pkg/plugin/implementation/keymanager"
 )
 
-// Mock KeyManager implementation
-type mockKeyManager struct{}
-
-func (m *mockKeyManager) SigningPublicKey(ctx context.Context, subscriberID, keyID string) (string, error) {
-	return "mock-signing-public-key", nil
-}
-
-func (m *mockKeyManager) SigningPrivateKey(ctx context.Context, subscriberID string) (string, string, error) {
-	return "mock-key-id", "mock-signing-private-key", nil
-}
-
-func (m *mockKeyManager) EncrPublicKey(ctx context.Context, subscriberID, keyID string) (string, error) {
-	return "mock-encryption-public-key", nil
-}
-
-func (m *mockKeyManager) EncrPrivateKey(ctx context.Context, subscriberID string) (string, string, error) {
-	return "mock-key-id", "mock-encryption-private-key", nil
-}
-
-func (m *mockKeyManager) DeletePrivateKeys(ctx context.Context, subscriberID string) error {
-	return nil
-}
-
-func (m *mockKeyManager) StorePrivateKeys(ctx context.Context, subscriberID string, keys *model.Keyset) error {
-	return nil
-}
-
-func (m *mockKeyManager) GenerateKeyPairs() (*model.Keyset, error) {
-	return &model.Keyset{
-		UniqueKeyID:    "mock-key-id",
-		SigningPrivate: "mock-signing-private-key",
-		SigningPublic:  "mock-signing-public-key",
-		EncrPrivate:    "mock-encryption-private-key",
-		EncrPublic:     "mock-encryption-public-key",
-	}, nil
-}
-
 type mockRegistry struct {
 	LookupFunc func(ctx context.Context, sub *model.Subscription) ([]model.Subscription, error)
 }
diff --git a/pkg/plugin/implementation/keymanager/keymanager_test.go b/pkg/plugin/implementation/keymanager/keymanager_test.go
index b08215e..f70e345 100644
--- a/pkg/plugin/implementation/keymanager/keymanager_test.go
+++ b/pkg/plugin/implementation/keymanager/keymanager_test.go
@@ -220,33 +220,6 @@ func TestGenerateKeyPairs(t *testing.T) {
 	}
 }
 
-type mockLogical struct {
-	writeFn func(path string, data map[string]interface{}) (*vault.Secret, error)
-}
-
-func (m *mockLogical) Write(path string, data map[string]interface{}) (*vault.Secret, error) {
-	return m.writeFn(path, data)
-}
-
-type mockClient struct {
-	*vault.Client
-	setTokenFn func(string)
-	logicalFn  func() *vault.Logical
-}
-
-func (m *mockClient) SetToken(token string) {
-	if m.setTokenFn != nil {
-		m.setTokenFn(token)
-	}
-}
-
-func (m *mockClient) Logical() *vault.Logical {
-	if m.logicalFn != nil {
-		return m.logicalFn()
-	}
-	return &vault.Logical{}
-}
-
 func TestGetVaultClient_Failures(t *testing.T) {
 	originalNewVaultClient := NewVaultClient
 	defer func() { NewVaultClient = originalNewVaultClient }()
@@ -297,7 +270,9 @@ func TestGetVaultClient_Failures(t *testing.T) {
 				return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 					w.WriteHeader(http.StatusOK)
 					w.Header().Set("Content-Type", "application/json")
-					io.WriteString(w, `{ "auth": null }`)
+					if _, err := io.WriteString(w, `{ "auth": null }`); err != nil {
+						t.Fatalf("failed to write response: %v", err)
+					}
 				}))
 			},
 			expectErr: "AppRole login failed: no auth info returned",
@@ -353,6 +328,13 @@ func TestGetVaultClient_Success(t *testing.T) {
 				"client_token": "mock-token"
 			}
 		}`)
+		if _, err := io.WriteString(w, `{
+			"auth": {
+				"client_token": "mock-token"
+			}
+		}`); err != nil {
+			t.Fatalf("failed to write response: %v", err)
+		}
 	}))
 	defer server.Close()
 
@@ -809,7 +791,9 @@ func setupMockVaultServer(t *testing.T, kvVersion, keyID string, success bool) *
 				}
 			}`, keyID)
 			w.Header().Set("Content-Type", "application/json")
-			w.Write([]byte(resp))
+			if _, err := w.Write([]byte(resp)); err != nil {
+				t.Fatalf("failed to write response: %v", err)
+			}
 		} else {
 			resp := fmt.Sprintf(`{
 				"data": {
@@ -821,7 +805,9 @@ func setupMockVaultServer(t *testing.T, kvVersion, keyID string, success bool) *
 				}
 			}`, keyID)
 			w.Header().Set("Content-Type", "application/json")
-			w.Write([]byte(resp))
+			if _, err := w.Write([]byte(resp)); err != nil {
+				t.Fatalf("failed to write response: %v", err)
+			}
 		}
 	})