steampipe/pkg/query/queryexecute/execute_test.go

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_NilInitData(t *testing.T) {
	ctx := context.Background()

	// This should not panic - function should validate initData is non-nil
	failures, err := RunBatchSession(ctx, nil)

	if err == nil {
		t.Fatal("Expected error when initData is nil, got nil")
	}

	if failures != 0 {
		t.Errorf("Expected 0 failures when initData is nil, got %d", failures)
	}
}

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")
	}
}

// TestRunBatchSession_LoadedTimeout demonstrates that RunBatchSession blocks forever
// if initData.Loaded never closes, even when the context is cancelled.
// References issue #4781
func TestRunBatchSession_LoadedTimeout(t *testing.T) {

	// Create a context with a short timeout
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()

	// Create InitData with a Loaded channel that will never close
	initData := &query.InitData{
		InitData: initialisation.InitData{
			Result: &db_common.InitResult{},
		},
		Loaded: make(chan struct{}), // This channel will never close
	}

	// This should return within the timeout, but currently blocks forever
	done := make(chan bool)
	var failures int
	var err error

	go func() {
		failures, err = RunBatchSession(ctx, initData)
		done <- true
	}()

	select {
	case <-done:
		// Function returned, check that it returned an error due to context cancellation
		assert.Error(t, err)
		assert.Equal(t, context.DeadlineExceeded, err)
		assert.Equal(t, 0, failures)
	case <-time.After(200 * time.Millisecond):
		t.Fatal("RunBatchSession blocked forever despite context cancellation - bug #4781")
	}
}

// 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)) {
}