package queryexecute import ( "context" "testing" "time" "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgxpool" "github.com/stretchr/testify/assert" "github.com/turbot/pipe-fittings/v2/modconfig" pqueryresult "github.com/turbot/pipe-fittings/v2/queryresult" "github.com/turbot/steampipe/v2/pkg/db/db_common" "github.com/turbot/steampipe/v2/pkg/export" "github.com/turbot/steampipe/v2/pkg/initialisation" "github.com/turbot/steampipe/v2/pkg/query" "github.com/turbot/steampipe/v2/pkg/query/queryresult" ) // Test Helpers // createMockInitData creates a mock InitData for testing func createMockInitData(t *testing.T) *query.InitData { t.Helper() initData := &query.InitData{ InitData: initialisation.InitData{ Result: &db_common.InitResult{}, ExportManager: export.NewManager(), Client: &mockClient{}, // Add mock client to prevent nil pointer panics }, Loaded: make(chan struct{}), StartTime: time.Now(), Queries: []*modconfig.ResolvedQuery{}, } return initData } // closeInitDataLoaded closes the Loaded channel to simulate initialization completion func closeInitDataLoaded(initData *query.InitData) { select { case <-initData.Loaded: // already closed default: close(initData.Loaded) } } // Test Suite: RunBatchSession func TestRunBatchSession_EmptyQueries(t *testing.T) { // ARRANGE: Create initData with no queries ctx := context.Background() initData := createMockInitData(t) initData.Queries = []*modconfig.ResolvedQuery{} // explicitly empty // Simulate successful initialization closeInitDataLoaded(initData) // ACT: Run batch session failures, err := RunBatchSession(ctx, initData) // ASSERT: Should return 0 failures and no error assert.NoError(t, err, "RunBatchSession should not error with empty queries") assert.Equal(t, 0, failures, "Should return 0 failures when no queries to execute") } func TestRunBatchSession_InitError(t *testing.T) { // ARRANGE: Create initData with an initialization error ctx := context.Background() initData := createMockInitData(t) // Simulate initialization error expectedErr := assert.AnError initData.Result.Error = expectedErr closeInitDataLoaded(initData) // ACT: Run batch session failures, err := RunBatchSession(ctx, initData) // ASSERT: Should return the init error immediately assert.Equal(t, expectedErr, err, "Should return initialization error") assert.Equal(t, 0, failures, "Should return 0 failures when init fails") } // TestRunBatchSession_NilClient tests that RunBatchSession handles nil Client gracefully func TestRunBatchSession_NilClient(t *testing.T) { // Create initData with nil Client initData := &query.InitData{ InitData: initialisation.InitData{ Result: &db_common.InitResult{}, Client: nil, // nil Client should be handled gracefully }, Loaded: make(chan struct{}), } // Signal that init is complete close(initData.Loaded) // This should not panic - it should handle nil Client gracefully _, err := RunBatchSession(context.Background(), initData) // We expect an error indicating that Client is required, not a panic if err == nil { t.Error("Expected error when Client is nil, got nil") } } // Test Suite: Helper Functions func TestNeedSnapshot_DefaultValues(t *testing.T) { // This test verifies the needSnapshot function behavior with default config // Note: This is a simple test but ensures the function doesn't panic // ACT: Call needSnapshot with default viper config result := needSnapshot() // ASSERT: Should return false with default settings assert.False(t, result, "needSnapshot should return false with default settings") } func TestShowBlankLineBetweenResults_DefaultValues(t *testing.T) { // This test verifies showBlankLineBetweenResults function with default config // ACT: Call function with default viper config result := showBlankLineBetweenResults() // ASSERT: Should return true with default settings (not CSV without header) assert.True(t, result, "Should show blank lines with default settings") } func TestHandlePublishSnapshotError_PaymentRequired(t *testing.T) { // ARRANGE: Create a 402 Payment Required error err := assert.AnError err = &mockError{msg: "402 Payment Required"} // ACT: Handle the error result := handlePublishSnapshotError(err) // ASSERT: Should reword the error message assert.Error(t, result) assert.Contains(t, result.Error(), "maximum number of snapshots reached") } func TestHandlePublishSnapshotError_OtherError(t *testing.T) { // ARRANGE: Create a different error err := assert.AnError // ACT: Handle the error result := handlePublishSnapshotError(err) // ASSERT: Should return the error unchanged assert.Equal(t, err, result) } // Test Suite: Edge Cases and Resource Management func TestExecuteQueries_EmptyQueriesList(t *testing.T) { // ARRANGE: InitData with empty queries list ctx := context.Background() initData := createMockInitData(t) initData.Queries = []*modconfig.ResolvedQuery{} // ACT: Execute queries directly failures := executeQueries(ctx, initData) // ASSERT: Should return 0 failures assert.Equal(t, 0, failures, "Should return 0 failures for empty queries list") } // Test Suite: Context and Cancellation func TestRunBatchSession_CancelHandlerSetup(t *testing.T) { // This test verifies that the cancel handler doesn't cause panics // We can't easily test the actual cancellation behavior without integration tests // ARRANGE ctx := context.Background() initData := createMockInitData(t) closeInitDataLoaded(initData) // ACT: Run batch session // Note: This test just verifies no panic occurs when setting up cancel handler assert.NotPanics(t, func() { _, _ = RunBatchSession(ctx, initData) }, "Should not panic when setting up cancel handler") } // Test Suite: Result Wrapping func TestWrapResult_NotNil(t *testing.T) { // This test ensures WrapResult doesn't panic and returns a valid wrapper // ARRANGE: Create a basic result from pipe-fittings // Note: We need to use the pipe-fittings queryresult package // This test verifies the wrapper functionality exists and doesn't panic wrapped := queryresult.NewResult(nil) // ASSERT: Should return a valid result assert.NotNil(t, wrapped, "NewResult should not return nil") } // Mock Types type mockError struct { msg string } func (e *mockError) Error() string { return e.msg } // mockClient is a minimal mock implementation of db_common.Client for testing type mockClient struct { customSearchPath []string requiredSearchPath []string } func (m *mockClient) Close(ctx context.Context) error { return nil } func (m *mockClient) LoadUserSearchPath(ctx context.Context) error { return nil } func (m *mockClient) SetRequiredSessionSearchPath(ctx context.Context) error { return nil } func (m *mockClient) GetRequiredSessionSearchPath() []string { return m.requiredSearchPath } func (m *mockClient) GetCustomSearchPath() []string { return m.customSearchPath } func (m *mockClient) AcquireManagementConnection(ctx context.Context) (*pgxpool.Conn, error) { return nil, nil } func (m *mockClient) AcquireSession(ctx context.Context) *db_common.AcquireSessionResult { return nil } func (m *mockClient) ExecuteSync(ctx context.Context, query string, args ...any) (*pqueryresult.SyncQueryResult, error) { return nil, nil } func (m *mockClient) Execute(ctx context.Context, query string, args ...any) (*queryresult.Result, error) { return nil, nil } func (m *mockClient) ExecuteSyncInSession(ctx context.Context, session *db_common.DatabaseSession, query string, args ...any) (*pqueryresult.SyncQueryResult, error) { return nil, nil } func (m *mockClient) ExecuteInSession(ctx context.Context, session *db_common.DatabaseSession, onConnectionLost func(), query string, args ...any) (*queryresult.Result, error) { return nil, nil } func (m *mockClient) ResetPools(ctx context.Context) { } func (m *mockClient) GetSchemaFromDB(ctx context.Context) (*db_common.SchemaMetadata, error) { return nil, nil } func (m *mockClient) ServerSettings() *db_common.ServerSettings { return nil } func (m *mockClient) RegisterNotificationListener(f func(notification *pgconn.Notification)) { }