add tests and update API to switch on type

Author: narsaynorath

static/app/views/explore/metrics/metricQuery.spec.tsx

@@ -1,10 +1,15 @@
+import {EQUATION_PREFIX} from 'sentry/utils/discover/fields';
 import {Mode} from 'sentry/views/explore/contexts/pageParamsContext/mode';
 import {
   decodeMetricsQueryParams,
+  defaultMetricQuery,
   encodeMetricQueryParams,
 } from 'sentry/views/explore/metrics/metricQuery';
 import {ReadableQueryParams} from 'sentry/views/explore/queryParams/readableQueryParams';
-import {VisualizeFunction} from 'sentry/views/explore/queryParams/visualize';
+import {
+  VisualizeEquation,
+  VisualizeFunction,
+} from 'sentry/views/explore/queryParams/visualize';
 
 describe('decodeMetricsQueryParams', () => {
   it('parses all visualizes', () => {
@@ -114,3 +119,41 @@ describe('decodeMetricsQueryParams', () => {
     );
   });
 });
+
+describe('defaultMetricQuery', () => {
+  it('returns a default metric query', () => {
+    const result = defaultMetricQuery();
+    expect(result).toEqual({
+      metric: {name: '', type: ''},
+      queryParams: new ReadableQueryParams({
+        extrapolate: true,
+        mode: Mode.SAMPLES,
+        query: '',
+        cursor: '',
+        fields: ['id', 'timestamp'],
+        sortBys: [{field: 'timestamp', kind: 'desc'}],
+        aggregateCursor: '',
+        aggregateFields: [new VisualizeFunction('sum(value)')],
+        aggregateSortBys: [{field: 'sum(value)', kind: 'desc'}],
+      }),
+    });
+  });
+
+  it('returns a default metric query with an equation', () => {
+    const result = defaultMetricQuery({type: 'equation'});
+    expect(result).toEqual({
+      metric: {name: '', type: ''},
+      queryParams: new ReadableQueryParams({
+        extrapolate: true,
+        mode: Mode.SAMPLES,
+        query: '',
+        cursor: '',
+        fields: ['id', 'timestamp'],
+        sortBys: [{field: 'timestamp', kind: 'desc'}],
+        aggregateCursor: '',
+        aggregateFields: [new VisualizeEquation(EQUATION_PREFIX)],
+        aggregateSortBys: [{field: EQUATION_PREFIX, kind: 'desc'}],
+      }),
+    });
+  });
+});

static/app/views/explore/metrics/metricQuery.tsx

@@ -105,8 +105,8 @@ export function encodeMetricQueryParams(metricQuery: BaseMetricQuery): string {
 }
 
 export function defaultMetricQuery({
-  equation,
-}: {equation?: boolean} = {}): BaseMetricQuery {
+  type = 'aggregate',
+}: {type?: 'aggregate' | 'equation'} = {}): BaseMetricQuery {
   return {
     metric: {name: '', type: ''},
     queryParams: new ReadableQueryParams({
@@ -119,9 +119,10 @@ export function defaultMetricQuery({
       sortBys: defaultSortBys(defaultFields()),
 
       aggregateCursor: '',
-      aggregateFields: equation ? [defaultAggregateEquation()] : defaultAggregateFields(),
+      aggregateFields:
+        type === 'equation' ? [defaultAggregateEquation()] : defaultAggregateFields(),
       aggregateSortBys: defaultAggregateSortBys(
-        equation ? [defaultAggregateEquation()] : defaultAggregateFields()
+        type === 'equation' ? [defaultAggregateEquation()] : defaultAggregateFields()
       ),
     }),
   };

static/app/views/explore/metrics/metricToolbar/aggregateDropdown.tsx

@@ -20,10 +20,7 @@ import {
   useSetMetricVisualizes,
 } from 'sentry/views/explore/metrics/metricsQueryParams';
 import {updateVisualizeYAxis} from 'sentry/views/explore/metrics/utils';
-import {
-  isVisualizeFunction,
-  VisualizeFunction,
-} from 'sentry/views/explore/queryParams/visualize';
+import {isVisualizeFunction} from 'sentry/views/explore/queryParams/visualize';
 
 const MULTI_SELECT_GROUP_KEYS = new Set(['percentiles', 'stats']);
 
@@ -36,6 +33,7 @@ export function AggregateDropdown({traceMetric}: {traceMetric: TraceMetric}) {
     return null;
   }
 
+  const narrowedVisualize = visualize;
   const groups = GROUPED_OPTIONS_BY_TYPE[traceMetric.type] ?? [];
   const selectedNames = new Set(
     visualizes.map(v => (isVisualizeFunction(v) ? (v.parsedFunction?.name ?? '') : ''))
@@ -45,15 +43,15 @@ export function AggregateDropdown({traceMetric}: {traceMetric: TraceMetric}) {
     if (selectedOptions.length === 0) {
       setMetricVisualizes([
         updateVisualizeYAxis(
-          visualize as VisualizeFunction,
+          narrowedVisualize,
           DEFAULT_YAXIS_BY_TYPE[traceMetric.type]!,
           traceMetric
         ),
       ]);
     } else {
       setMetricVisualizes(
         selectedOptions.map(o =>
-          updateVisualizeYAxis(visualize as VisualizeFunction, o.value, traceMetric)
+          updateVisualizeYAxis(narrowedVisualize, o.value, traceMetric)
         )
       );
     }

static/app/views/explore/metrics/metricsTab.spec.tsx

@@ -1,3 +1,4 @@
+import {OrganizationFixture} from 'sentry-fixture/organization';
 import {TimeSeriesFixture} from 'sentry-fixture/timeSeries';
 import {
   createTraceMetricFixtures,
@@ -582,6 +583,35 @@ describe('MetricsTabContent', () => {
     });
     expect(parsedQuery.aggregateFields).toContainEqual({groupBy: 'test.region'});
   });
+
+  it('does not show the Add Equation button when the feature flag is disabled', async () => {
+    render(
+      <ProviderWrapper>
+        <MetricsTabContent datePageFilterProps={datePageFilterProps} />
+      </ProviderWrapper>,
+      {
+        organization,
+      }
+    );
+    expect(await screen.findByText('Add Metric')).toBeInTheDocument();
+    expect(screen.queryByText('Add Equation')).not.toBeInTheDocument();
+  });
+
+  it('shows the Add Equation button when the feature flag is enabled', async () => {
+    const orgWithFeature = OrganizationFixture({
+      features: ['tracemetrics-enabled', 'tracemetrics-equations-in-explore'],
+    });
+    render(
+      <ProviderWrapper>
+        <MetricsTabContent datePageFilterProps={datePageFilterProps} />
+      </ProviderWrapper>,
+      {
+        organization: orgWithFeature,
+      }
+    );
+    expect(await screen.findByText('Add Metric')).toBeInTheDocument();
+    expect(screen.getByText('Add Equation')).toBeInTheDocument();
+  });
 });
 
 describe('MetricsTabContent (tracemetrics-ui-refresh)', () => {

static/app/views/explore/metrics/metricsTab.tsx

@@ -34,7 +34,6 @@ import {MetricToolbar} from 'sentry/views/explore/metrics/metricToolbar';
 import {MetricSaveAs} from 'sentry/views/explore/metrics/metricToolbar/metricSaveAs';
 import {
   MultiMetricsQueryParamsProvider,
-  useAddEquationQuery,
   useAddMetricQuery,
   useMultiMetricsQueryParams,
 } from 'sentry/views/explore/metrics/multiMetricsQueryParams';
@@ -121,17 +120,19 @@ function MetricsQueryBuilderSection({
   const organization = useOrganization();
   const metricQueries = useMultiMetricsQueryParams();
   const addMetricQuery = useAddMetricQuery();
-  const addEquationQuery = useAddEquationQuery();
+  const addEquationQuery = useAddMetricQuery({type: 'equation'});
   const hasEquations = useHasMetricEquations();
   const references = useMemo(() => {
     return new Set(
       metricQueries
-        .filter(metricQuery =>
+        .map((metricQuery, queryIndex) => ({metricQuery, queryIndex}))
+        .filter(({metricQuery}) =>
           metricQuery.queryParams.visualizes.some(isVisualizeFunction)
         )
-        .map((_metricQuery, index) => getVisualizeLabel(index))
+        .map(({queryIndex}) => getVisualizeLabel(queryIndex))
     );
   }, [metricQueries]);
+
   if (canUseMetricsUIRefresh(organization)) {
     return (
       <ExploreControlSection expanded={controlSectionExpanded ?? true}>