diff --git a/Features/Onboarding/Sources/Navigation/ImportWalletNavigationStack.swift b/Features/Onboarding/Sources/Navigation/ImportWalletNavigationStack.swift
index 80b7d2d42..5bb1e7fd0 100644
--- a/Features/Onboarding/Sources/Navigation/ImportWalletNavigationStack.swift
+++ b/Features/Onboarding/Sources/Navigation/ImportWalletNavigationStack.swift
@@ -2,7 +2,6 @@
 
 import SwiftUI
 import Primitives
-import Localization
 import PrimitivesComponents
 
 public struct ImportWalletNavigationStack: View {
@@ -82,24 +81,11 @@ extension ImportWalletNavigationStack {
         }
     }
 
-    func onImportComplete(data: WalletImportData) {
-        Task {
-            do {
-                let wallet = try await model.importWallet(data: data)
-                navigate(to: .walletProfile(wallet: wallet))
-            } catch {
-                debugLog("Failed to import wallet: \(error)")
-            }
-        }
+    func onImportComplete(wallet: Wallet) {
+        navigate(to: .walletProfile(wallet: wallet))
     }
 
     func onSetupWalletComplete(_ wallet: Wallet) {
-        Task {
-            do {
-                try await model.setupWalletComplete()
-            } catch {
-                debugLog("Failed to setup wallet: \(error)")
-            }
-        }
+        model.onComplete?()
     }
 }
diff --git a/Features/Onboarding/Sources/ViewModels/ImportWalletSceneViewModel.swift b/Features/Onboarding/Sources/ViewModels/ImportWalletSceneViewModel.swift
index d74637cf1..181364b62 100644
--- a/Features/Onboarding/Sources/ViewModels/ImportWalletSceneViewModel.swift
+++ b/Features/Onboarding/Sources/ViewModels/ImportWalletSceneViewModel.swift
@@ -29,13 +29,13 @@ final class ImportWalletSceneViewModel {
     var isPresentingScanner = false
     var isPresentingAlertMessage: AlertMessage?
 
-    private let onComplete: ((WalletImportData) -> Void)?
+    private let onComplete: ((Wallet) -> Void)?
 
     init(
         walletService: WalletService,
         nameService: any NameServiceable,
         type: ImportWalletType,
-        onComplete: ((WalletImportData) -> Void)?
+        onComplete: ((Wallet) -> Void)?
     ) {
         self.walletService = walletService
         self.nameService = nameService
@@ -175,12 +175,12 @@ extension ImportWalletSceneViewModel {
             }
             switch type {
             case .multicoin:
-                importWallet(
+                try await importWallet(
                     name: recipient.name,
                     keystoreType: .phrase(words: words, chains: AssetConfiguration.allChains)
                 )
             case .chain(let chain):
-                importWallet(
+                try await importWallet(
                     name: recipient.name,
                     keystoreType: .single(words: words, chain: chain)
                 )
@@ -189,7 +189,7 @@ extension ImportWalletSceneViewModel {
             guard try validateForm(type: importType, address: recipient.address, words: [trimmedInput]) else {
                 return
             }
-            importWallet(name: recipient.name, keystoreType: .privateKey(text: trimmedInput, chain: chain!))
+            try await importWallet(name: recipient.name, keystoreType: .privateKey(text: trimmedInput, chain: chain!))
         case .address:
             guard try validateForm(type: importType, address: recipient.address, words: []) else {
                 return
@@ -197,12 +197,16 @@ extension ImportWalletSceneViewModel {
             let chain = chain!
             let address = chain.checksumAddress(recipient.address)
 
-            importWallet(name: recipient.name, keystoreType: .address(address: address, chain: chain))
+            try await importWallet(name: recipient.name, keystoreType: .address(address: address, chain: chain))
         }
     }
 
-    private func importWallet(name: String, keystoreType: KeystoreImportType) {
-        onComplete?(WalletImportData(name: name, keystoreType: keystoreType))
+    private func importWallet(name: String, keystoreType: KeystoreImportType) async throws {
+        let wallet = try await walletService.loadOrCreateWallet(name: name, type: keystoreType, source: .import)
+        walletService.acceptTerms()
+        try await walletService.setCurrent(wallet: wallet)
+        buttonState = .normal
+        onComplete?(wallet)
     }
 
     private func validateForm(type: WalletImportType, address: String, words: [String]) throws  -> Bool {
diff --git a/Features/Onboarding/Sources/ViewModels/ImportWalletViewModel.swift b/Features/Onboarding/Sources/ViewModels/ImportWalletViewModel.swift
index b3d745e9f..d51b8bb02 100644
--- a/Features/Onboarding/Sources/ViewModels/ImportWalletViewModel.swift
+++ b/Features/Onboarding/Sources/ViewModels/ImportWalletViewModel.swift
@@ -34,10 +34,6 @@ public final class ImportWalletViewModel {
     public var isAcceptTermsCompleted: Bool {
         walletService.isAcceptTermsCompleted
     }
-
-    func dismiss() {
-        onComplete?()
-    }
 }
 
 // MARK: - Actions
@@ -46,15 +42,4 @@ extension ImportWalletViewModel {
     func presentSelectImage(wallet: Wallet) {
         isPresentingSelectImageWallet = wallet
     }
-
-    func importWallet(data: WalletImportData) async throws -> Wallet {
-        let wallet = try await walletService.loadOrCreateWallet(name: data.name, type: data.keystoreType, source: .import)
-        walletService.acceptTerms()
-        try await walletService.setCurrent(wallet: wallet)
-        return wallet
-    }
-
-    func setupWalletComplete() async throws {
-        dismiss()
-    }
 }
diff --git a/Features/Onboarding/Tests/ImportWalletViewModelTests.swift b/Features/Onboarding/Tests/ImportWalletViewModelTests.swift
deleted file mode 100644
index 82f7f3163..000000000
--- a/Features/Onboarding/Tests/ImportWalletViewModelTests.swift
+++ /dev/null
@@ -1,40 +0,0 @@
-// Copyright (c). Gem Wallet. All rights reserved.
-
-import Foundation
-import Testing
-import Primitives
-import Preferences
-import Keystore
-import KeystoreTestKit
-import PrimitivesTestKit
-import NameServiceTestKit
-import StoreTestKit
-import WalletServiceTestKit
-
-@testable import Onboarding
-
-@MainActor
-struct ImportWalletViewModelTests {
-
-    @Test
-    func importWalletDoesNotSetAddressStatus() async throws {
-        let keystore = KeystoreMock()
-        let preferences = WalletPreferences(walletId: Wallet.mock().walletId)
-        preferences.clear()
-
-        let model = ImportWalletViewModel(
-            walletService: .mock(keystore: keystore),
-            avatarService: .init(store: .mock()),
-            nameService: .mock(),
-            onComplete: nil
-        )
-
-        _ = try await model.importWallet(data: WalletImportData(
-            name: "Test",
-            keystoreType: .phrase(words: LocalKeystore.words, chains: [.tron])
-        ))
-
-        #expect(preferences.completeInitialAddressStatus == false)
-        #expect(preferences.completeInitialLoadAssets == false)
-    }
-}