▽

Source file src/github.com/cybertec-postgresql/pgwatch/v6/internal/metrics/yaml_test.go

Documentation: github.com/cybertec-postgresql/pgwatch/v6/internal/metrics

     1  package metrics_test
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"path/filepath"
     7  	"runtime"
     8  	"sync"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
    13  	"github.com/stretchr/testify/assert"
    14  	"gopkg.in/yaml.v3"
    15  )
    16  
    17  func TestDeaultMetrics(t *testing.T) {
    18  	fmr, err := metrics.NewYAMLMetricReaderWriter(ctx, "") // empty path is reserved for default metrics
    19  	assert.NoError(t, err)
    20  
    21  	// Test GetMetrics
    22  	metricsDefs, err := fmr.GetMetrics()
    23  	assert.NoError(t, err)
    24  
    25  	// Test WriteMetrics
    26  	err = fmr.WriteMetrics(metricsDefs)
    27  	assert.Error(t, err)
    28  
    29  	// Test DeleteMetric
    30  	err = fmr.DeleteMetric("test")
    31  	assert.Error(t, err)
    32  
    33  	// Test UpdateMetric
    34  	err = fmr.UpdateMetric("test", metrics.Metric{})
    35  	assert.Error(t, err)
    36  
    37  	// Test DeletePreset
    38  	err = fmr.DeletePreset("test")
    39  	assert.Error(t, err)
    40  
    41  	// Test UpdatePreset
    42  	err = fmr.UpdatePreset("test", metrics.Preset{})
    43  	assert.Error(t, err)
    44  }
    45  
    46  func TestWriteMetricsToFile(t *testing.T) {
    47  	// Define test data
    48  	metricDefs := metrics.Metrics{
    49  		MetricDefs: metrics.MetricDefs{
    50  			"test_metric": metrics.Metric{
    51  				SQLs: map[int]string{
    52  					1: "SELECT 1",
    53  				},
    54  				InitSQL:         "SELECT 1",
    55  				NodeStatus:      "primary",
    56  				Gauges:          []string{"gauge1", "gauge2"},
    57  				IsInstanceLevel: true,
    58  				StorageName:     "storage1",
    59  				Description:     "Test metric",
    60  			},
    61  		},
    62  		PresetDefs: metrics.PresetDefs{
    63  			"test_preset": metrics.Preset{
    64  				Description: "Test preset",
    65  				Metrics: metrics.MetricIntervals{
    66  					"test_metric": 1.0,
    67  				},
    68  			},
    69  		},
    70  	}
    71  
    72  	// Create a temporary file for testing
    73  	tempDir := t.TempDir()
    74  	tempFile := filepath.Join(tempDir, "metrics.yaml")
    75  
    76  	fmr, err := metrics.NewYAMLMetricReaderWriter(ctx, tempFile)
    77  	assert.NoError(t, err)
    78  
    79  	// Call the function being tested
    80  	err = fmr.WriteMetrics(&metricDefs)
    81  	assert.NoError(t, err)
    82  
    83  	// Read the contents of the file
    84  	metrics, err := fmr.GetMetrics()
    85  	assert.NoError(t, err)
    86  
    87  	// Assert that the file contains the expected data
    88  	assert.Equal(t, metricDefs, *metrics)
    89  }
    90  func TestMetricsToFile(t *testing.T) {
    91  	// Define test data
    92  	metricDefs := metrics.Metrics{
    93  		MetricDefs: metrics.MetricDefs{
    94  			"existing_metric": metrics.Metric{
    95  				SQLs: map[int]string{
    96  					1: "SELECT 1",
    97  				},
    98  				InitSQL:         "SELECT 1",
    99  				NodeStatus:      "primary",
   100  				Gauges:          []string{"gauge1", "gauge2"},
   101  				IsInstanceLevel: true,
   102  				StorageName:     "storage1",
   103  				Description:     "Existing metric",
   104  			},
   105  		},
   106  		PresetDefs: metrics.PresetDefs{
   107  			"test_preset": metrics.Preset{
   108  				Description: "Test preset",
   109  				Metrics: metrics.MetricIntervals{
   110  					"existing_metric": 1.0,
   111  				},
   112  			},
   113  		},
   114  	}
   115  
   116  	// Create a temporary file for testing
   117  	tempDir := t.TempDir()
   118  	tempFile := filepath.Join(tempDir, "metrics.yaml")
   119  
   120  	fmr, err := metrics.NewYAMLMetricReaderWriter(ctx, tempFile)
   121  	assert.NoError(t, err)
   122  
   123  	// Write initial metrics to the file
   124  	err = fmr.WriteMetrics(&metricDefs)
   125  	assert.NoError(t, err)
   126  
   127  	// Call the function being tested
   128  	newMetric := metrics.Metric{
   129  		SQLs: map[int]string{
   130  			1: "SELECT 2",
   131  		},
   132  		InitSQL:         "SELECT 2",
   133  		NodeStatus:      "primary",
   134  		Gauges:          []string{"gauge3", "gauge4"},
   135  		IsInstanceLevel: true,
   136  		StorageName:     "storage2",
   137  		Description:     "New metric",
   138  	}
   139  	err = fmr.UpdateMetric("new_metric", newMetric)
   140  	assert.NoError(t, err)
   141  
   142  	// Read the updated metrics from the file
   143  	updatedMetrics, err := fmr.GetMetrics()
   144  	assert.NoError(t, err)
   145  
   146  	// Assert that the metric was updated correctly
   147  	expectedMetrics := metricDefs
   148  	expectedMetrics.MetricDefs["new_metric"] = newMetric
   149  	assert.Equal(t, expectedMetrics, *updatedMetrics)
   150  
   151  	// Call the function being tested
   152  	err = fmr.DeleteMetric("new_metric")
   153  	assert.NoError(t, err)
   154  
   155  	// Read the updated metrics from the file
   156  	updatedMetrics, err = fmr.GetMetrics()
   157  	assert.NoError(t, err)
   158  
   159  	// Assert that the metric was deleted correctly
   160  	assert.Zero(t, updatedMetrics.MetricDefs["new_metric"])
   161  }
   162  
   163  func TestPresetsToFile(t *testing.T) {
   164  	// Define test data
   165  	presetDefs := metrics.PresetDefs{
   166  		"existing_preset": metrics.Preset{
   167  			Description: "Existing preset",
   168  			Metrics: metrics.MetricIntervals{
   169  				"existing_metric": 1.0,
   170  			},
   171  		},
   172  	}
   173  
   174  	// Create a temporary file for testing
   175  	tempDir := t.TempDir()
   176  	tempFile := filepath.Join(tempDir, "metrics.yaml")
   177  
   178  	fmr, err := metrics.NewYAMLMetricReaderWriter(ctx, tempFile)
   179  	assert.NoError(t, err)
   180  
   181  	// Write initial presets to the file
   182  	err = fmr.WriteMetrics(&metrics.Metrics{
   183  		PresetDefs: presetDefs,
   184  	})
   185  	assert.NoError(t, err)
   186  
   187  	// Call the function being tested
   188  	newPreset := metrics.Preset{
   189  		Description: "New preset",
   190  		Metrics: metrics.MetricIntervals{
   191  			"new_metric": 1.0,
   192  		},
   193  	}
   194  	err = fmr.UpdatePreset("new_preset", newPreset)
   195  	assert.NoError(t, err)
   196  
   197  	// Read the updated presets from the file
   198  	updatedMetrics, err := fmr.GetMetrics()
   199  	assert.NoError(t, err)
   200  
   201  	// Assert that the preset was updated correctly
   202  	expectedPresets := presetDefs
   203  	expectedPresets["new_preset"] = newPreset
   204  	assert.Equal(t, expectedPresets, updatedMetrics.PresetDefs)
   205  
   206  	// check the delete preset function
   207  	err = fmr.DeletePreset("new_preset")
   208  	assert.NoError(t, err)
   209  
   210  	// Read the updated presets from the file
   211  	updatedMetrics, err = fmr.GetMetrics()
   212  	assert.NoError(t, err)
   213  
   214  	// Assert that the preset was deleted correctly
   215  	assert.Zero(t, updatedMetrics.PresetDefs["new_preset"])
   216  }
   217  
   218  func TestErrorHandlingToFile(t *testing.T) {
   219  	fmr, err := metrics.NewYAMLMetricReaderWriter(ctx, "/") // empty path is reserved for default metrics
   220  	assert.NoError(t, err)
   221  
   222  	// Test WriteMetrics
   223  	err = fmr.WriteMetrics(&metrics.Metrics{})
   224  	assert.Error(t, err)
   225  
   226  	// Test GetMetrics - root dir has no yaml files, returns empty metrics without error
   227  	_, err = fmr.GetMetrics()
   228  	assert.NoError(t, err)
   229  
   230  	// Test DeleteMetric
   231  	err = fmr.DeleteMetric("test")
   232  	assert.Error(t, err)
   233  
   234  	// Test UpdateMetric
   235  	err = fmr.UpdateMetric("test", metrics.Metric{})
   236  	assert.Error(t, err)
   237  
   238  	// Test DeletePreset
   239  	err = fmr.DeletePreset("test")
   240  	assert.Error(t, err)
   241  
   242  	// Test UpdatePreset
   243  	err = fmr.UpdatePreset("test", metrics.Preset{})
   244  	assert.Error(t, err)
   245  
   246  	// Test invalid YAML
   247  	tempDir := t.TempDir()
   248  	tempFile := filepath.Join(tempDir, "metrics.yaml")
   249  	file, err := os.Create(tempFile)
   250  	assert.NoError(t, err)
   251  	defer file.Close()
   252  	_, err = file.WriteString("invalid yaml")
   253  	assert.NoError(t, err)
   254  
   255  	fmr, err = metrics.NewYAMLMetricReaderWriter(ctx, tempFile)
   256  	assert.NoError(t, err)
   257  
   258  	_, err = fmr.GetMetrics()
   259  	assert.Error(t, err)
   260  }
   261  
   262  func TestCreateMetricAndPreset(t *testing.T) {
   263  
   264  	a := assert.New(t)
   265  
   266  	t.Run("YAML_CreateMetric_Success", func(t *testing.T) {
   267  		tmpFile := filepath.Join(t.TempDir(), "test_metrics.yaml")
   268  		defer os.Remove(tmpFile)
   269  
   270  		// Create YAML reader/writer
   271  		yamlrw, err := metrics.NewYAMLMetricReaderWriter(ctx, tmpFile)
   272  		a.NoError(err)
   273  
   274  		// Initialize empty metrics file first
   275  		emptyMetrics := &metrics.Metrics{
   276  			MetricDefs: make(map[string]metrics.Metric),
   277  			PresetDefs: make(map[string]metrics.Preset),
   278  		}
   279  		err = yamlrw.WriteMetrics(emptyMetrics)
   280  		a.NoError(err)
   281  
   282  		// Create a new metric
   283  		testMetric := metrics.Metric{
   284  			Description: "Test metric for creation",
   285  		}
   286  		err = yamlrw.CreateMetric("test_metric", testMetric)
   287  		a.NoError(err)
   288  
   289  		// Verify it was created
   290  		m, err := yamlrw.GetMetrics()
   291  		a.NoError(err)
   292  		a.Contains(m.MetricDefs, "test_metric")
   293  		a.Equal("Test metric for creation", m.MetricDefs["test_metric"].Description)
   294  
   295  		// Try to create the same metric again - should fail
   296  		err = yamlrw.CreateMetric("test_metric", testMetric)
   297  		a.Error(err)
   298  		a.Equal(metrics.ErrMetricExists, err)
   299  	})
   300  
   301  	t.Run("YAML_CreatePreset_Success", func(t *testing.T) {
   302  		tmpFile := filepath.Join(t.TempDir(), "test_presets.yaml")
   303  		defer os.Remove(tmpFile)
   304  
   305  		yamlrw, err := metrics.NewYAMLMetricReaderWriter(ctx, tmpFile)
   306  		a.NoError(err)
   307  
   308  		// Initialize empty metrics file first
   309  		emptyMetrics := &metrics.Metrics{
   310  			MetricDefs: make(map[string]metrics.Metric),
   311  			PresetDefs: make(map[string]metrics.Preset),
   312  		}
   313  		err = yamlrw.WriteMetrics(emptyMetrics)
   314  		a.NoError(err)
   315  
   316  		// Create a new preset
   317  		testPreset := metrics.Preset{
   318  			Description: "Test preset for creation",
   319  			Metrics:     metrics.MetricIntervals{"db_stats": 60},
   320  		}
   321  		err = yamlrw.CreatePreset("test_preset", testPreset)
   322  		a.NoError(err)
   323  
   324  		// Verify it was created
   325  		m, err := yamlrw.GetMetrics()
   326  		a.NoError(err)
   327  		a.Contains(m.PresetDefs, "test_preset")
   328  		a.Equal("Test preset for creation", m.PresetDefs["test_preset"].Description)
   329  
   330  		// Try to create the same preset again - should fail
   331  		err = yamlrw.CreatePreset("test_preset", testPreset)
   332  		a.Error(err)
   333  		a.Equal(metrics.ErrPresetExists, err)
   334  	})
   335  }
   336  
   337  func TestMetricsDir(t *testing.T) {
   338  	a := assert.New(t)
   339  
   340  	// first metrics file data
   341  	metrics1 := metrics.Metrics{
   342  		MetricDefs: map[string]metrics.Metric{
   343  			"metric1": {
   344  				Description: "metric1 description",
   345  			},
   346  		},
   347  		PresetDefs: map[string]metrics.Preset{
   348  			"preset1": {
   349  				Description: "preset1 description",
   350  				Metrics: metrics.MetricIntervals{
   351  					"metric1": 10,
   352  				},
   353  			},
   354  		},
   355  	}
   356  
   357  	// second metrics file data
   358  	metrics2 := metrics.Metrics{
   359  		MetricDefs: map[string]metrics.Metric{
   360  			"metric2": {
   361  				Description: "metric2 description",
   362  			},
   363  		},
   364  		PresetDefs: map[string]metrics.Preset{
   365  			"preset2": {
   366  				Description: "preset2 description",
   367  				Metrics: metrics.MetricIntervals{
   368  					"metric2": 10,
   369  				},
   370  			},
   371  		},
   372  	}
   373  
   374  	metrics1File, err := yaml.Marshal(metrics1)
   375  	a.NoError(err)
   376  	metrics2File, err := yaml.Marshal(metrics2)
   377  	a.NoError(err)
   378  
   379  	// write data to different files in a folder: one .yaml and one .yml
   380  	tempDir := t.TempDir()
   381  	err = os.WriteFile(filepath.Join(tempDir, "metrics1.yaml"), metrics1File, 0644)
   382  	a.NoError(err)
   383  	err = os.WriteFile(filepath.Join(tempDir, "metrics2.yml"), metrics2File, 0644)
   384  	a.NoError(err)
   385  
   386  	// use folder of yaml files for metrics configs
   387  	yamlrw, err := metrics.NewYAMLMetricReaderWriter(ctx, tempDir)
   388  	a.NoError(err)
   389  
   390  	// load metrics configs from folder
   391  	ms, err := yamlrw.GetMetrics()
   392  	a.NoError(err)
   393  	a.Equal("metric1 description", ms.MetricDefs["metric1"].Description)
   394  	a.Equal("preset1 description", ms.PresetDefs["preset1"].Description)
   395  	a.Equal("metric2 description", ms.MetricDefs["metric2"].Description)
   396  	a.Equal("preset2 description", ms.PresetDefs["preset2"].Description)
   397  }
   398  
   399  func TestConcurrentMetricUpdates(t *testing.T) {
   400  	a := assert.New(t)
   401  	tempDir := t.TempDir()
   402  	tempFile := filepath.Join(tempDir, "metrics.yaml")
   403  
   404  	yamlrw, err := metrics.NewYAMLMetricReaderWriter(ctx, tempFile)
   405  	a.NoError(err)
   406  
   407  	// Create initial empty metrics file
   408  	initialMetrics := &metrics.Metrics{
   409  		MetricDefs: make(map[string]metrics.Metric),
   410  		PresetDefs: make(map[string]metrics.Preset),
   411  	}
   412  	err = yamlrw.WriteMetrics(initialMetrics)
   413  	a.NoError(err)
   414  
   415  	numGoroutines := 10
   416  	var wg sync.WaitGroup
   417  
   418  	// Each goroutine will add a unique metric
   419  	for id := range numGoroutines {
   420  		wg.Go(func() {
   421  			metricName := fmt.Sprintf("metric_%d", id)
   422  			testMetric := metrics.Metric{
   423  				Description: fmt.Sprintf("Test metric %d", id),
   424  				SQLs: map[int]string{
   425  					1: fmt.Sprintf("SELECT %d", id),
   426  				},
   427  			}
   428  			time.Sleep(time.Millisecond * time.Duration(id%3))
   429  			err := yamlrw.UpdateMetric(metricName, testMetric)
   430  			a.NoError(err, "Error during concurrent update")
   431  		})
   432  	}
   433  	wg.Wait()
   434  
   435  	finalMetrics, err := yamlrw.GetMetrics()
   436  	a.NoError(err)
   437  	a.Equal(numGoroutines, len(finalMetrics.MetricDefs), "Some updates were lost due to race condition!")
   438  }
   439  
   440  func TestConcurrentPresetUpdates(t *testing.T) {
   441  	a := assert.New(t)
   442  	tempDir := t.TempDir()
   443  	tempFile := filepath.Join(tempDir, "metrics.yaml")
   444  
   445  	yamlrw, err := metrics.NewYAMLMetricReaderWriter(ctx, tempFile)
   446  	a.NoError(err)
   447  
   448  	// Create initial empty metrics file
   449  	initialMetrics := &metrics.Metrics{
   450  		MetricDefs: make(map[string]metrics.Metric),
   451  		PresetDefs: make(map[string]metrics.Preset),
   452  	}
   453  	err = yamlrw.WriteMetrics(initialMetrics)
   454  	a.NoError(err)
   455  
   456  	numGoroutines := 10
   457  	var wg sync.WaitGroup
   458  
   459  	for id := range numGoroutines {
   460  		wg.Go(func() {
   461  			presetName := fmt.Sprintf("preset_%d", id)
   462  			testPreset := metrics.Preset{
   463  				Description: fmt.Sprintf("Test preset %d", id),
   464  				Metrics:     metrics.MetricIntervals{fmt.Sprintf("metric_%d", id): 60},
   465  			}
   466  			time.Sleep(time.Millisecond * time.Duration(id%3))
   467  			err := yamlrw.UpdatePreset(presetName, testPreset)
   468  			a.NoError(err, "Error during concurrent update")
   469  		})
   470  	}
   471  	wg.Wait()
   472  
   473  	// ensure all presets were saved
   474  	finalMetrics, err := yamlrw.GetMetrics()
   475  	a.NoError(err)
   476  	a.Equal(numGoroutines, len(finalMetrics.PresetDefs), "Some updates were lost due to race condition!")
   477  }
   478  
   479  func TestGetMetricsNonExistentPath(t *testing.T) {
   480  	nonExistent := filepath.Join(t.TempDir(), "does_not_exist", "metrics.yaml")
   481  	fmr, err := metrics.NewYAMLMetricReaderWriter(ctx, nonExistent)
   482  	assert.NoError(t, err)
   483  	_, err = fmr.GetMetrics()
   484  	assert.Error(t, err)
   485  }
   486  
   487  func TestGetMetricsDirWithInvalidYAML(t *testing.T) {
   488  	dir := t.TempDir()
   489  	err := os.WriteFile(filepath.Join(dir, "bad.yaml"), []byte("invalid: yaml: {unclosed"), 0644)
   490  	assert.NoError(t, err)
   491  	fmr, err := metrics.NewYAMLMetricReaderWriter(ctx, dir)
   492  	assert.NoError(t, err)
   493  	_, err = fmr.GetMetrics()
   494  	assert.Error(t, err)
   495  }
   496  
   497  func TestMutationsGetMetricsError(t *testing.T) {
   498  	nonExistent := filepath.Join(t.TempDir(), "does_not_exist", "metrics.yaml")
   499  	fmr, err := metrics.NewYAMLMetricReaderWriter(ctx, nonExistent)
   500  	assert.NoError(t, err)
   501  
   502  	assert.Error(t, fmr.DeleteMetric("x"))
   503  	assert.Error(t, fmr.UpdateMetric("x", metrics.Metric{}))
   504  	assert.Error(t, fmr.CreateMetric("x", metrics.Metric{}))
   505  	assert.Error(t, fmr.DeletePreset("x"))
   506  	assert.Error(t, fmr.UpdatePreset("x", metrics.Preset{}))
   507  	assert.Error(t, fmr.CreatePreset("x", metrics.Preset{}))
   508  }
   509  
   510  func TestLoadMetricsUnreadableFile(t *testing.T) {
   511  	if runtime.GOOS == "windows" {
   512  		t.Skip("cannot reliably test file permissions on Windows")
   513  	}
   514  	if os.Getuid() == 0 {
   515  		t.Skip("running as root, permission checks do not apply")
   516  	}
   517  	f := filepath.Join(t.TempDir(), "metrics.yaml")
   518  	err := os.WriteFile(f, []byte(""), 0000)
   519  	assert.NoError(t, err)
   520  	fmr, err := metrics.NewYAMLMetricReaderWriter(ctx, f)
   521  	assert.NoError(t, err)
   522  	_, err = fmr.GetMetrics()
   523  	assert.Error(t, err)
   524  }
   525