github.com/Uptycs/basequery-go@v0.8.0/client_test.go (about)

     1  package osquery
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  
     8  	"github.com/Uptycs/basequery-go/gen/osquery"
     9  	"github.com/Uptycs/basequery-go/mock"
    10  	"github.com/stretchr/testify/assert"
    11  )
    12  
    13  func TestQueryRows(t *testing.T) {
    14  	mock := &mock.ExtensionManager{}
    15  	client := &ExtensionManagerClient{Client: mock}
    16  
    17  	// Transport related error
    18  	mock.QueryFunc = func(ctx context.Context, sql string) (*osquery.ExtensionResponse, error) {
    19  		return nil, errors.New("Boom")
    20  	}
    21  	rows, err := client.QueryRows("select 1")
    22  	assert.NotNil(t, err)
    23  	row, err := client.QueryRow("select 1")
    24  	assert.NotNil(t, err)
    25  
    26  	// Nil status
    27  	mock.QueryFunc = func(ctx context.Context, sql string) (*osquery.ExtensionResponse, error) {
    28  		return &osquery.ExtensionResponse{}, nil
    29  	}
    30  	rows, err = client.QueryRows("select 1")
    31  	assert.NotNil(t, err)
    32  	row, err = client.QueryRow("select 1")
    33  	assert.NotNil(t, err)
    34  
    35  	// Query error
    36  	mock.QueryFunc = func(ctx context.Context, sql string) (*osquery.ExtensionResponse, error) {
    37  		return &osquery.ExtensionResponse{
    38  			Status: &osquery.ExtensionStatus{Code: 1, Message: "bad query"},
    39  		}, nil
    40  	}
    41  	rows, err = client.QueryRows("select bad query")
    42  	assert.NotNil(t, err)
    43  	row, err = client.QueryRow("select bad query")
    44  	assert.NotNil(t, err)
    45  
    46  	// Good query (one row)
    47  	expectedRows := []map[string]string{
    48  		{"1": "1"},
    49  	}
    50  	mock.QueryFunc = func(ctx context.Context, sql string) (*osquery.ExtensionResponse, error) {
    51  		return &osquery.ExtensionResponse{
    52  			Status:   &osquery.ExtensionStatus{Code: 0, Message: "OK"},
    53  			Response: expectedRows,
    54  		}, nil
    55  	}
    56  	rows, err = client.QueryRows("select 1")
    57  	assert.Nil(t, err)
    58  	assert.Equal(t, expectedRows, rows)
    59  	row, err = client.QueryRow("select 1")
    60  	assert.Nil(t, err)
    61  	assert.Equal(t, expectedRows[0], row)
    62  
    63  	// Good query (multiple rows)
    64  	expectedRows = []map[string]string{
    65  		{"1": "1"},
    66  		{"1": "2"},
    67  	}
    68  	mock.QueryFunc = func(ctx context.Context, sql string) (*osquery.ExtensionResponse, error) {
    69  		return &osquery.ExtensionResponse{
    70  			Status:   &osquery.ExtensionStatus{Code: 0, Message: "OK"},
    71  			Response: expectedRows,
    72  		}, nil
    73  	}
    74  	rows, err = client.QueryRows("select 1 union select 2")
    75  	assert.Nil(t, err)
    76  	assert.Equal(t, expectedRows, rows)
    77  	row, err = client.QueryRow("select 1 union select 2")
    78  	assert.NotNil(t, err)
    79  }