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 }