shardmapper tests for string->string

pkg/logql/shardmapper_test.go

@@ -119,6 +119,51 @@ func TestMapSampleExpr(t *testing.T) {
 	}
 }
 
+func TestMappingStrings(t *testing.T) {
+	m, err := NewShardMapper(2)
+	require.Nil(t, err)
+	for _, tc := range []struct {
+		in  string
+		out string
+	}{
+		{
+			in:  `sum(rate({foo="bar"}[1m]))`,
+			out: `sum(downstream<sum(rate(({foo="bar"})[1m])), shard=0_of_2> ++ downstream<sum(rate(({foo="bar"})[1m])), shard=1_of_2>)`,
+		},
+		{
+			in:  `max(count(rate({foo="bar"}[5m]))) / 2`,
+			out: `max(sum(downstream<count(rate(({foo="bar"})[5m])), shard=0_of_2> ++ downstream<count(rate(({foo="bar"})[5m])), shard=1_of_2>)) / 2.000000`,
+		},
+		{
+			in:  `topk(3, rate({foo="bar"}[5m]))`,
+			out: `topk(3,downstream<rate(({foo="bar"})[5m]), shard=0_of_2> ++ downstream<rate(({foo="bar"})[5m]), shard=1_of_2>)`,
+		},
+		{
+			in:  `sum(max(rate({foo="bar"}[5m])))`,
+			out: `sum(max(downstream<rate(({foo="bar"})[5m]), shard=0_of_2> ++ downstream<rate(({foo="bar"})[5m]), shard=1_of_2>))`,
+		},
+		{
+			in:  `{foo="bar"} |= "id=123"`,
+			out: `downstream<{foo="bar"}|="id=123", shard=0_of_2> ++ downstream<{foo="bar"}|="id=123", shard=1_of_2>`,
+		},
+		{
+			in:  `sum by (cluster) (rate({foo="bar"} |= "id=123" [5m]))`,
+			out: `sum by(cluster)(downstream<sum by(cluster)(rate(({foo="bar"}|="id=123")[5m])), shard=0_of_2> ++ downstream<sum by(cluster)(rate(({foo="bar"}|="id=123")[5m])), shard=1_of_2>)`,
+		},
+	} {
+		t.Run(tc.in, func(t *testing.T) {
+			ast, err := ParseExpr(tc.in)
+			require.Nil(t, err)
+
+			mapped, err := m.Map(ast)
+			require.Nil(t, err)
+
+			require.Equal(t, tc.out, mapped.String())
+
+		})
+	}
+}
+
 func TestMapping(t *testing.T) {
 	m, err := NewShardMapper(2)
 	require.Nil(t, err)