import { describe, it, expect, vi, beforeEach } from "vitest"; import { render, screen, fireEvent, act, waitFor } from "@testing-library/react"; import { useDbConnection } from "./useDbConnection"; import { useDbViewerStore } from "../stores/dbViewerStore"; import * as commands from "../lib/commands"; vi.mock("../lib/commands", () => ({ dbConnect: vi.fn().mockResolvedValue(undefined), dbDisconnect: vi.fn().mockResolvedValue(undefined), getDatabases: vi.fn().mockResolvedValue([]), getSchemas: vi.fn().mockResolvedValue([]), getTables: vi.fn().mockResolvedValue([]), getConnectionPassword: vi.fn().mockResolvedValue("pw"), getConnectionSshPassword: vi.fn().mockResolvedValue(null), getConnectionSshPassphrase: vi.fn().mockResolvedValue(null), })); const mockCommands = vi.mocked(commands); const mockConnection = { id: "c1", name: "Conn", db_type: "postgresql", host: "localhost", port: 5432, username: "postgres", database: "mydb", folder_id: null, keychain_ref: null, tag_ids: [], environment: null, ssh_host: null as string | null, ssh_port: null, ssh_user: null, ssh_auth_method: null as string | null, ssh_private_key_path: null, ssl_mode: null, ssl_ca_path: null, ssl_cert_path: null, ssl_key_path: null, use_keychain: undefined as boolean | undefined, created_at: "2026-07-26T00:00:00Z", updated_at: "2026-07-26T00:00:00Z", }; // Mutable stand-in for the connection store's session-password behavior: a // fallback password (default "pw", like a keychain hit) plus a per-connection // session map written by setSessionPassword (like the real store). const mockStoreState = vi.hoisted(() => { const sessionPasswords = new Map(); let fallbackPassword: string | null = "pw"; return { sessionPasswords, setFallbackPassword: (p: string | null) => { fallbackPassword = p; }, getConnectionPassword: vi.fn( async (id: string) => sessionPasswords.get(id) ?? fallbackPassword, ), setSessionPassword: vi.fn((id: string, pw: string) => { sessionPasswords.set(id, pw); }), }; }); vi.mock("../stores/connectionStore", () => ({ useConnectionStore: { getState: () => ({ connections: [mockConnection], getConnectionPassword: mockStoreState.getConnectionPassword, setSessionPassword: mockStoreState.setSessionPassword, }), }, })); function Harness() { const { connect } = useDbConnection("c1"); return ( ); } function PromptHarness() { const { passwordPromptOpen } = useDbConnection("c1"); return
{passwordPromptOpen ? "open" : "closed"}
; } describe("useDbConnection", () => { beforeEach(() => { vi.clearAllMocks(); // Reset the mutable store stand-in + connection for isolation. mockStoreState.sessionPasswords.clear(); mockStoreState.setFallbackPassword("pw"); (mockConnection as any).use_keychain = undefined; useDbViewerStore.getState().reset(); mockCommands.getDatabases.mockResolvedValue(["mydb", "otherdb"]); mockCommands.getSchemas.mockResolvedValue(["app", "public"]); mockCommands.getTables.mockResolvedValue([]); mockCommands.getConnectionSshPassword.mockResolvedValue(null); mockCommands.getConnectionSshPassphrase.mockResolvedValue(null); mockConnection.ssh_host = null; mockConnection.ssh_auth_method = null; }); it("sets schemaTreeLoading around the connect fetch", async () => { mockCommands.getSchemas.mockImplementation( () => new Promise((res) => setTimeout(() => res(["public"]), 50)), ); render(); fireEvent.click(screen.getByText("connect")); await waitFor(() => expect(useDbViewerStore.getState().schemaTreeLoading).toBe(true), ); await waitFor(() => expect(useDbViewerStore.getState().schemaTreeLoading).toBe(false), ); }); it("clears schemaTreeLoading when connect throws", async () => { mockCommands.getSchemas.mockRejectedValue(new Error("boom")); render(); fireEvent.click(screen.getByText("connect")); await waitFor(() => expect(useDbViewerStore.getState().schemaTreeLoading).toBe(false), ); }); it("connects and smart-selects the public schema when available", async () => { render(); fireEvent.click(screen.getByText("connect")); await waitFor(() => { expect(useDbViewerStore.getState().currentDatabase).toBe("mydb"); }); expect(useDbViewerStore.getState().currentSchema).toBe("public"); // tables are fetched for the smart default schema expect(mockCommands.getTables).toHaveBeenCalledWith("c1", "public"); }); it("falls back to the first schema when public is absent", async () => { mockCommands.getSchemas.mockResolvedValue(["analytics", "app"]); render(); fireEvent.click(screen.getByText("connect")); await waitFor(() => { expect(useDbViewerStore.getState().currentSchema).toBe("analytics"); }); expect(mockCommands.getTables).toHaveBeenCalledWith("c1", "analytics"); }); it("reconnects and refreshes schemas when the database changes", async () => { render(); fireEvent.click(screen.getByText("connect")); await waitFor(() => { expect(useDbViewerStore.getState().currentDatabase).toBe("mydb"); }); mockCommands.getDatabases.mockResolvedValue(["mydb", "otherdb"]); mockCommands.getSchemas.mockResolvedValue(["analytics", "public"]); await act(async () => { useDbViewerStore.setState({ currentDatabase: "otherdb" }); }); expect(mockCommands.dbConnect).toHaveBeenLastCalledWith( "c1", expect.objectContaining({ database: "otherdb" }), ); expect(useDbViewerStore.getState().schemas).toEqual([ "analytics", "public", ]); expect(useDbViewerStore.getState().currentSchema).toBe("public"); // tables refetched for the new database's smart schema expect(mockCommands.getTables).toHaveBeenLastCalledWith("c1", "public"); }); it("refreshes schemas when switching back to the initial database", async () => { render(); fireEvent.click(screen.getByText("connect")); await waitFor(() => { expect(useDbViewerStore.getState().currentDatabase).toBe("mydb"); }); // Switch away to otherdb await act(async () => { useDbViewerStore.setState({ currentDatabase: "otherdb" }); }); await waitFor(() => expect(mockCommands.dbConnect).toHaveBeenLastCalledWith( "c1", expect.objectContaining({ database: "otherdb" }), ), ); // Switch back to mydb — must reconnect and refresh schemas again mockCommands.getSchemas.mockResolvedValue(["public", "reporting"]); await act(async () => { useDbViewerStore.setState({ currentDatabase: "mydb" }); }); expect(mockCommands.dbConnect).toHaveBeenLastCalledWith( "c1", expect.objectContaining({ database: "mydb" }), ); expect(useDbViewerStore.getState().schemas).toEqual([ "public", "reporting", ]); expect(useDbViewerStore.getState().currentSchema).toBe("public"); }); it("opens the password prompt when use_keychain=false and no session password", async () => { mockConnection.use_keychain = false; mockStoreState.setFallbackPassword(null); render(); await waitFor(() => expect(screen.getByTestId("prompt")).toHaveTextContent("open"), ); // Early return: no Tauri call fires before the user submits a password. expect(mockCommands.dbConnect).not.toHaveBeenCalled(); }); it("does not prompt when keychain is enabled (even without a stored password)", async () => { mockConnection.use_keychain = true; mockStoreState.setFallbackPassword(null); render(); await waitFor(() => expect(mockCommands.dbConnect).toHaveBeenCalled(), ); expect(screen.getByTestId("prompt")).toHaveTextContent("closed"); }); it("submitPassword stores the session password and reconnects", async () => { mockConnection.use_keychain = false; mockStoreState.setFallbackPassword(null); let submit: ((pw: string) => void) | null = null; function SubmitHarness() { const { passwordPromptOpen, submitPassword } = useDbConnection("c1"); submit = submitPassword; return (
{passwordPromptOpen ? "open" : "closed"}
); } render(); await waitFor(() => expect(screen.getByTestId("prompt")).toHaveTextContent("open"), ); act(() => submit!("secret")); expect(mockStoreState.setSessionPassword).toHaveBeenCalledWith("c1", "secret"); await waitFor(() => expect(screen.getByTestId("prompt")).toHaveTextContent("closed"), ); await waitFor(() => expect(mockCommands.dbConnect).toHaveBeenCalled()); }); it("fetches ssh secrets from keychain before connecting when ssh_host is set", async () => { mockConnection.ssh_host = "bastion.example.com"; mockConnection.ssh_auth_method = "password"; mockCommands.getConnectionSshPassword.mockResolvedValue("sshpw"); mockCommands.getConnectionSshPassphrase.mockResolvedValue(null); render(); fireEvent.click(screen.getByText("connect")); await waitFor(() => { expect(mockCommands.dbConnect).toHaveBeenCalled(); }); const config = mockCommands.dbConnect.mock.calls[0][1]; expect(mockCommands.getConnectionSshPassword).toHaveBeenCalledWith("c1"); expect(mockCommands.getConnectionSshPassphrase).toHaveBeenCalledWith("c1"); expect(config.ssh_password).toBe("sshpw"); expect(config.ssh_passphrase).toBeNull(); }); });