github.com/kubewharf/katalyst-core@v0.5.3/pkg/controller/resource-recommend/datasource/datasource_proxy_test.go (about)

     1  /*
     2  Copyright 2022 The Katalyst Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package datasource
    18  
    19  import (
    20  	"errors"
    21  	"testing"
    22  	"time"
    23  
    24  	"github.com/stretchr/testify/assert"
    25  	"github.com/stretchr/testify/mock"
    26  
    27  	datasourcetypes "github.com/kubewharf/katalyst-core/pkg/util/resource-recommend/types/datasource"
    28  )
    29  
    30  type MockDatasource struct {
    31  	mock.Mock
    32  }
    33  
    34  func (m *MockDatasource) ConvertMetricToQuery(metric datasourcetypes.Metric) (*datasourcetypes.Query, error) {
    35  	args := m.Called(metric)
    36  	return args.Get(0).(*datasourcetypes.Query), args.Error(1)
    37  }
    38  
    39  func (m *MockDatasource) QueryTimeSeries(query *datasourcetypes.Query, start, end time.Time, step time.Duration) (*datasourcetypes.TimeSeries, error) {
    40  	args := m.Called(query, start, end, step)
    41  	return args.Get(0).(*datasourcetypes.TimeSeries), args.Error(1)
    42  }
    43  
    44  func TestProxy_QueryTimeSeries(t *testing.T) {
    45  	// Create mock objects
    46  	mockDatasource := &MockDatasource{}
    47  	mockTimeSeries := &datasourcetypes.TimeSeries{}
    48  
    49  	// Create Proxy with mock datasource
    50  	proxy := &Proxy{
    51  		datasourceMap: map[DatasourceType]Datasource{
    52  			DatasourceType("mock"): mockDatasource,
    53  		},
    54  	}
    55  
    56  	// Define test cases
    57  	testCases := []struct {
    58  		name        string
    59  		datasource  DatasourceType
    60  		metric      datasourcetypes.Metric
    61  		start       time.Time
    62  		end         time.Time
    63  		step        time.Duration
    64  		expectedTS  *datasourcetypes.TimeSeries
    65  		expectedErr error
    66  	}{
    67  		{
    68  			name:       "Valid query",
    69  			datasource: "mock",
    70  			metric:     datasourcetypes.Metric{},
    71  			start:      time.Now(),
    72  			end:        time.Now(),
    73  			step:       time.Second,
    74  			expectedTS: mockTimeSeries,
    75  		},
    76  		{
    77  			name:        "Error getting datasource",
    78  			datasource:  "invalid",
    79  			metric:      datasourcetypes.Metric{},
    80  			start:       time.Now(),
    81  			end:         time.Now(),
    82  			step:        time.Second,
    83  			expectedErr: errors.New("datasource not found"),
    84  		},
    85  	}
    86  
    87  	// Set up mock behaviors and perform tests
    88  	for _, tc := range testCases {
    89  		t.Run(tc.name, func(t *testing.T) {
    90  			// Set up mock behaviors
    91  			mockDatasource.On("ConvertMetricToQuery", tc.metric).Return(&datasourcetypes.Query{}, nil)
    92  			if tc.expectedErr == nil {
    93  				mockDatasource.On("QueryTimeSeries", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockTimeSeries, nil)
    94  			} else {
    95  				mockDatasource.On("QueryTimeSeries", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, tc.expectedErr)
    96  			}
    97  
    98  			// Call the QueryTimeSeries method
    99  			result, err := proxy.QueryTimeSeries(tc.datasource, tc.metric, tc.start, tc.end, tc.step)
   100  
   101  			// Verify the results
   102  			assert.Equal(t, tc.expectedTS, result)
   103  			assert.Equal(t, tc.expectedErr, err)
   104  
   105  			// Verify mock invocations
   106  			mockDatasource.AssertExpectations(t)
   107  		})
   108  	}
   109  }