github.com/google/cloudprober@v0.11.3/probes/probes_test.go (about)

     1  // Copyright 2017 The Cloudprober Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package probes_test
    16  
    17  import (
    18  	"context"
    19  	"testing"
    20  
    21  	"github.com/golang/protobuf/proto"
    22  	"github.com/google/cloudprober/metrics"
    23  	"github.com/google/cloudprober/probes"
    24  	"github.com/google/cloudprober/probes/options"
    25  	configpb "github.com/google/cloudprober/probes/proto"
    26  	testdatapb "github.com/google/cloudprober/probes/testdata"
    27  	targetspb "github.com/google/cloudprober/targets/proto"
    28  )
    29  
    30  var testProbeIntialized int
    31  
    32  type testProbe struct{}
    33  
    34  func (p *testProbe) Init(name string, opts *options.Options) error {
    35  	testProbeIntialized++
    36  	return nil
    37  }
    38  
    39  func (p *testProbe) Start(ctx context.Context, dataChan chan *metrics.EventMetrics) {}
    40  
    41  func TestGetExtensionProbe(t *testing.T) {
    42  	probeDef := &configpb.ProbeDef{
    43  		Name: proto.String("ext-probe"),
    44  		Type: configpb.ProbeDef_EXTENSION.Enum(),
    45  		Targets: &targetspb.TargetsDef{
    46  			Type: &targetspb.TargetsDef_DummyTargets{},
    47  		},
    48  	}
    49  
    50  	// This has the same effect as using the following in your config:
    51  	// probe {
    52  	//    name: "ext-probe"
    53  	//    targets: ...
    54  	//    ...
    55  	//    [cloudprober.probes.testdata.fancy_probe] {
    56  	//      name: "fancy"
    57  	//    }
    58  	// }
    59  	err := proto.SetExtension(probeDef, testdatapb.E_FancyProbe, &testdatapb.FancyProbe{Name: proto.String("fancy")})
    60  	if err != nil {
    61  		t.Fatalf("error setting up extension in test probe proto: %v", err)
    62  	}
    63  	probeInfo, err := probes.CreateProbe(probeDef, &options.Options{})
    64  	if err == nil {
    65  		t.Errorf("Expected error in building probe from extensions, got probe: %v", probeInfo)
    66  	}
    67  	t.Log(err.Error())
    68  
    69  	// Register our test probe type and try again.
    70  	probes.RegisterProbeType(200, func() probes.Probe {
    71  		return &testProbe{}
    72  	})
    73  
    74  	probeInfo, err = probes.CreateProbe(probeDef, &options.Options{})
    75  	if err != nil {
    76  		t.Errorf("Got error in building probe from extensions: %v", err)
    77  	}
    78  	if probeInfo == nil || probeInfo.Name != "ext-probe" {
    79  		t.Errorf("Extension probe not in the probes map")
    80  	}
    81  	_, ok := probeInfo.Probe.(*testProbe)
    82  	if !ok {
    83  		t.Errorf("Extension probe (%v) is not of type *testProbe", probeInfo)
    84  	}
    85  	if testProbeIntialized != 1 {
    86  		t.Errorf("Extensions probe's Init() called %d times, should be called exactly once.", testProbeIntialized)
    87  	}
    88  }