github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/querier/astmapper/subtree_folder_test.go (about)

     1  package astmapper
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  
     7  	"github.com/pkg/errors"
     8  	"github.com/prometheus/prometheus/promql/parser"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestPredicate(t *testing.T) {
    13  	for i, tc := range []struct {
    14  		input    string
    15  		fn       predicate
    16  		expected bool
    17  		err      bool
    18  	}{
    19  		{
    20  			input: "selector1{} or selector2{}",
    21  			fn: predicate(func(node parser.Node) (bool, error) {
    22  				return false, errors.New("some err")
    23  			}),
    24  			expected: false,
    25  			err:      true,
    26  		},
    27  		{
    28  			input: "selector1{} or selector2{}",
    29  			fn: predicate(func(node parser.Node) (bool, error) {
    30  				return false, nil
    31  			}),
    32  			expected: false,
    33  			err:      false,
    34  		},
    35  		{
    36  			input: "selector1{} or selector2{}",
    37  			fn: predicate(func(node parser.Node) (bool, error) {
    38  				return true, nil
    39  			}),
    40  			expected: true,
    41  			err:      false,
    42  		},
    43  		{
    44  			input:    `sum without(__cortex_shard__) (__embedded_queries__{__cortex_queries__="tstquery"}) or sum(selector)`,
    45  			fn:       predicate(isEmbedded),
    46  			expected: true,
    47  			err:      false,
    48  		},
    49  	} {
    50  		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
    51  			expr, err := parser.ParseExpr(tc.input)
    52  			require.Nil(t, err)
    53  
    54  			res, err := Predicate(expr.(parser.Node), tc.fn)
    55  			if tc.err {
    56  				require.Error(t, err)
    57  			} else {
    58  				require.Nil(t, err)
    59  			}
    60  
    61  			require.Equal(t, tc.expected, res)
    62  		})
    63  	}
    64  }
    65  
    66  func TestSubtreeMapper(t *testing.T) {
    67  	for i, tc := range []struct {
    68  		input    string
    69  		expected string
    70  	}{
    71  		// embed an entire histogram
    72  		{
    73  			input:    "histogram_quantile(0.5, rate(alertmanager_http_request_duration_seconds_bucket[1m]))",
    74  			expected: `__embedded_queries__{__cortex_queries__="{\"Concat\":[\"histogram_quantile(0.5, rate(alertmanager_http_request_duration_seconds_bucket[1m]))\"]}"}`,
    75  		},
    76  		// embed a binary expression across two functions
    77  		{
    78  			input:    `rate(http_requests_total{cluster="eu-west2"}[5m]) or rate(http_requests_total{cluster="us-central1"}[5m])`,
    79  			expected: `__embedded_queries__{__cortex_queries__="{\"Concat\":[\"rate(http_requests_total{cluster=\\\"eu-west2\\\"}[5m]) or rate(http_requests_total{cluster=\\\"us-central1\\\"}[5m])\"]}"}`,
    80  		},
    81  
    82  		// the first leg (histogram) hasn't been embedded at any level, so embed that, but ignore the right leg
    83  		// which has already been embedded.
    84  		{
    85  			input: `sum(histogram_quantile(0.5, rate(selector[1m]))) +
    86  				sum without(__cortex_shard__) (__embedded_queries__{__cortex_queries__="tstquery"})`,
    87  			expected: `
    88  			  __embedded_queries__{__cortex_queries__="{\"Concat\":[\"sum(histogram_quantile(0.5, rate(selector[1m])))\"]}"} +
    89  			  sum without(__cortex_shard__) (__embedded_queries__{__cortex_queries__="tstquery"})
    90  `,
    91  		},
    92  		// should not embed scalars
    93  		{
    94  			input:    `histogram_quantile(0.5, __embedded_queries__{__cortex_queries__="tstquery"})`,
    95  			expected: `histogram_quantile(0.5, __embedded_queries__{__cortex_queries__="tstquery"})`,
    96  		},
    97  	} {
    98  		t.Run(fmt.Sprintf("[%d]", i), func(t *testing.T) {
    99  			mapper := NewSubtreeFolder()
   100  
   101  			expr, err := parser.ParseExpr(tc.input)
   102  			require.Nil(t, err)
   103  			res, err := mapper.Map(expr)
   104  			require.Nil(t, err)
   105  
   106  			expected, err := parser.ParseExpr(tc.expected)
   107  			require.Nil(t, err)
   108  
   109  			require.Equal(t, expected.String(), res.String())
   110  
   111  		})
   112  	}
   113  }