cloud-ip-validator/internal/db/queries_registry_levels_test.go

package db

import (
	"reflect"
	"testing"
)

func TestCheckLevelAndFamily(t *testing.T) {
	for source, want := range map[string]string{
		"egress": LevelEgress, "inbound-site-1": LevelIngress, "inbound-site-12": LevelIngress,
		"": "", "other": "",
	} {
		if got := CheckLevel(source); got != want {
			t.Errorf("CheckLevel(%q) = %q, want %q", source, got, want)
		}
	}
	for ct, want := range map[string]string{
		"https": "https", "icmp": "icmp", "ssh": "ssh",
		"tcp-22": "tcp", "tcp-443": "tcp", "tls-443": "tls", "dns": "dns", "-x": "-x",
	} {
		if got := CheckFamily(ct); got != want {
			t.Errorf("CheckFamily(%q) = %q, want %q", ct, got, want)
		}
	}
}

// addCheck records one check for the address's current queue row.
func addCheck(t *testing.T, d *DB, addr, source, checkType, target string, success bool) {
	t.Helper()
	ctx := t.Context()
	ip, err := d.GetIPByAddress(ctx, addr)
	if err != nil {
		t.Fatalf("get %s: %v", addr, err)
	}
	if err := d.UpsertCheck(ctx, Check{
		IPID: ip.ID, IPAddress: addr, AttemptNumber: ip.AttemptNumber,
		Source: source, CheckType: checkType, Target: target,
		Success: success, CheckedAt: Now(),
	}); err != nil {
		t.Fatalf("upsert check: %v", err)
	}
}

func TestRegistryLevelsGroupByTypeAndLevel(t *testing.T) {
	d, ctx := newTestDB(t)
	if _, err := d.SubmitIPs(ctx, []string{"1.2.3.4"}); err != nil {
		t.Fatal(err)
	}
	a := "1.2.3.4"
	// Egress: https 2 of 3, icmp 1 of 1.
	addCheck(t, d, a, SourceEgress, "https", "https://a.test", true)
	addCheck(t, d, a, SourceEgress, "https", "https://b.test", true)
	addCheck(t, d, a, SourceEgress, "https", "https://c.test", false)
	addCheck(t, d, a, SourceEgress, "icmp", "a.test", true)
	// Ingress from two sites: tcp-22 and tcp-443 are one family; tls, ssh,
	// icmp and a type unknown today ("dns") are listed on their own.
	for site := 1; site <= 2; site++ {
		src := InboundSource(site)
		addCheck(t, d, a, src, "tcp-22", a, true)
		addCheck(t, d, a, src, "tcp-443", a, site == 1)
		addCheck(t, d, a, src, "tls-443", a, true)
		addCheck(t, d, a, src, "ssh", a, false)
		addCheck(t, d, a, src, "icmp", a, true)
		addCheck(t, d, a, src, "dns", a, true)
	}
	// A source that is neither egress nor an inbound site is not counted.
	addCheck(t, d, a, "manual", "https", "x", true)

	s, err := d.GetRegistryByAddress(ctx, a)
	if err != nil {
		t.Fatal(err)
	}
	wantEgress := LevelResult{Total: 4, OK: 3, ByType: []TypeStat{
		{"https", 3, 2}, {"icmp", 1, 1},
	}}
	wantIngress := LevelResult{Total: 12, OK: 9, ByType: []TypeStat{
		{"dns", 2, 2}, {"icmp", 2, 2}, {"ssh", 2, 0}, {"tcp", 4, 3}, {"tls", 2, 2},
	}}
	if !reflect.DeepEqual(s.Egress, wantEgress) {
		t.Errorf("egress = %+v, want %+v", s.Egress, wantEgress)
	}
	if !reflect.DeepEqual(s.Ingress, wantIngress) {
		t.Errorf("ingress = %+v, want %+v", s.Ingress, wantIngress)
	}
	if s.LastCycleID != 1 {
		t.Errorf("LastCycleID = %d, want 1", s.LastCycleID)
	}
}

func TestRegistryLevelsNoChecksAreZero(t *testing.T) {
	d, ctx := newTestDB(t)
	if _, err := d.SubmitIPs(ctx, []string{"1.2.3.4"}); err != nil {
		t.Fatal(err)
	}
	s, err := d.GetRegistryByAddress(ctx, "1.2.3.4")
	if err != nil {
		t.Fatal(err)
	}
	if s.LastCycleID != 0 || s.Egress.Total != 0 || s.Ingress.Total != 0 || len(s.Egress.ByType) != 0 {
		t.Fatalf("expected empty levels, got %+v", s)
	}
}

// The counts follow the newest cycle, also after the queue row is deleted,
// and one grouped query serves several addresses without mixing them up.
func TestRegistryLevelsLatestCycleAndManyAddresses(t *testing.T) {
	d, ctx := newTestDB(t)
	if _, err := d.SubmitIPs(ctx, []string{"1.1.1.1", "2.2.2.2"}); err != nil {
		t.Fatal(err)
	}
	addCheck(t, d, "1.1.1.1", SourceEgress, "https", "t1", false)
	addCheck(t, d, "2.2.2.2", SourceEgress, "https", "t1", true)
	addCheck(t, d, "2.2.2.2", InboundSource(1), "tcp-22", "2.2.2.2", true)

	// New cycle for 1.1.1.1: delete and submit again.
	ip, err := d.GetIPByAddress(ctx, "1.1.1.1")
	if err != nil {
		t.Fatal(err)
	}
	if err := d.DeleteIP(ctx, ip.ID); err != nil {
		t.Fatal(err)
	}
	if _, err := d.SubmitIPs(ctx, []string{"1.1.1.1"}); err != nil {
		t.Fatal(err)
	}
	addCheck(t, d, "1.1.1.1", SourceEgress, "icmp", "t2", true)
	addCheck(t, d, "1.1.1.1", SourceEgress, "https", "t2", true)

	all, err := d.ListRegistry(ctx)
	if err != nil {
		t.Fatal(err)
	}
	got := map[string]RegistrySummary{}
	for _, s := range all {
		got[s.IPAddress] = s
	}
	a := got["1.1.1.1"]
	if a.LastCycleID != 2 || a.Egress.Total != 2 || a.Egress.OK != 2 || a.Ingress.Total != 0 {
		t.Errorf("1.1.1.1: %+v", a)
	}
	b := got["2.2.2.2"]
	if b.LastCycleID != 1 || b.Egress.Total != 1 || b.Ingress.Total != 1 || b.Ingress.OK != 1 {
		t.Errorf("2.2.2.2: %+v", b)
	}

	// Page and single lookups agree with the full list.
	page, _, err := d.ListRegistryPage(ctx, RegistryFilter{}, 10, 0)
	if err != nil {
		t.Fatal(err)
	}
	for _, s := range page {
		if !reflect.DeepEqual(s.Egress, got[s.IPAddress].Egress) || !reflect.DeepEqual(s.Ingress, got[s.IPAddress].Ingress) {
			t.Errorf("page differs from list for %s", s.IPAddress)
		}
	}

	// Delete the queue row of 2.2.2.2: the counts stay (history outlives it).
	ip2, _ := d.GetIPByAddress(ctx, "2.2.2.2")
	if err := d.DeleteIP(ctx, ip2.ID); err != nil {
		t.Fatal(err)
	}
	s, err := d.GetRegistryByAddress(ctx, "2.2.2.2")
	if err != nil {
		t.Fatal(err)
	}
	if s.Egress.Total != 1 || s.Ingress.Total != 1 {
		t.Errorf("after delete: %+v", s)
	}
}