feat: refactor

internal/dataset/dataset_test.go

@@ -0,0 +1,152 @@
+package dataset
+
+import (
+	"testing"
+)
+
+func TestDatasetShape(t *testing.T) {
+	if NumHours != 65 {
+		t.Errorf("NumHours = %d, want 65", NumHours)
+	}
+	if NumLevels != 47 {
+		t.Errorf("NumLevels = %d, want 47", NumLevels)
+	}
+	if NumVariables != 3 {
+		t.Errorf("NumVariables = %d, want 3", NumVariables)
+	}
+	if NumLatitudes != 361 {
+		t.Errorf("NumLatitudes = %d, want 361", NumLatitudes)
+	}
+	if NumLongitudes != 720 {
+		t.Errorf("NumLongitudes = %d, want 720", NumLongitudes)
+	}
+}
+
+func TestDatasetSize(t *testing.T) {
+	// 65 * 47 * 3 * 361 * 720 * 4 = 9,528,667,200
+	want := int64(9_528_667_200)
+	if DatasetSize != want {
+		t.Errorf("DatasetSize = %d, want %d", DatasetSize, want)
+	}
+}
+
+func TestPressureLevels(t *testing.T) {
+	if len(Pressures) != 47 {
+		t.Fatalf("len(Pressures) = %d, want 47", len(Pressures))
+	}
+
+	// First should be 1000 (highest pressure, near surface)
+	if Pressures[0] != 1000 {
+		t.Errorf("Pressures[0] = %d, want 1000", Pressures[0])
+	}
+	// Last should be 1 (lowest pressure, high atmosphere)
+	if Pressures[46] != 1 {
+		t.Errorf("Pressures[46] = %d, want 1", Pressures[46])
+	}
+
+	// Should be sorted descending
+	for i := 1; i < len(Pressures); i++ {
+		if Pressures[i] >= Pressures[i-1] {
+			t.Errorf("Pressures not descending at [%d]: %d >= %d", i, Pressures[i], Pressures[i-1])
+		}
+	}
+
+	// Total levels: 26 from pgrb2 + 21 from pgrb2b = 47
+	if len(PressuresPgrb2) != 26 {
+		t.Errorf("len(PressuresPgrb2) = %d, want 26", len(PressuresPgrb2))
+	}
+	if len(PressuresPgrb2b) != 21 {
+		t.Errorf("len(PressuresPgrb2b) = %d, want 21", len(PressuresPgrb2b))
+	}
+}
+
+func TestPressureIndex(t *testing.T) {
+	if PressureIndex(1000) != 0 {
+		t.Errorf("PressureIndex(1000) = %d, want 0", PressureIndex(1000))
+	}
+	if PressureIndex(1) != 46 {
+		t.Errorf("PressureIndex(1) = %d, want 46", PressureIndex(1))
+	}
+	if PressureIndex(500) != 20 {
+		t.Errorf("PressureIndex(500) = %d, want 20", PressureIndex(500))
+	}
+	if PressureIndex(9999) != -1 {
+		t.Errorf("PressureIndex(9999) = %d, want -1", PressureIndex(9999))
+	}
+}
+
+func TestPressureLevelSet(t *testing.T) {
+	// 1000 mb should be in pgrb2 (A)
+	ls, ok := PressureLevelSet(1000)
+	if !ok || ls != LevelSetA {
+		t.Errorf("PressureLevelSet(1000) = %v, %v; want A, true", ls, ok)
+	}
+
+	// 125 mb should be in pgrb2b (B)
+	ls, ok = PressureLevelSet(125)
+	if !ok || ls != LevelSetB {
+		t.Errorf("PressureLevelSet(125) = %v, %v; want B, true", ls, ok)
+	}
+
+	// 1, 2, 3, 5, 7 should be in pgrb2b (B)
+	for _, p := range []int{1, 2, 3, 5, 7} {
+		ls, ok := PressureLevelSet(p)
+		if !ok || ls != LevelSetB {
+			t.Errorf("PressureLevelSet(%d) = %v, %v; want B, true", p, ls, ok)
+		}
+	}
+
+	// Every pressure level should have a level set assignment
+	for _, p := range Pressures {
+		_, ok := PressureLevelSet(p)
+		if !ok {
+			t.Errorf("PressureLevelSet(%d) not found", p)
+		}
+	}
+}
+
+func TestHourIndex(t *testing.T) {
+	if HourIndex(0) != 0 {
+		t.Errorf("HourIndex(0) = %d, want 0", HourIndex(0))
+	}
+	if HourIndex(3) != 1 {
+		t.Errorf("HourIndex(3) = %d, want 1", HourIndex(3))
+	}
+	if HourIndex(192) != 64 {
+		t.Errorf("HourIndex(192) = %d, want 64", HourIndex(192))
+	}
+	if HourIndex(1) != -1 {
+		t.Errorf("HourIndex(1) = %d, want -1 (not multiple of 3)", HourIndex(1))
+	}
+	if HourIndex(195) != -1 {
+		t.Errorf("HourIndex(195) = %d, want -1 (out of range)", HourIndex(195))
+	}
+}
+
+func TestHours(t *testing.T) {
+	hours := Hours()
+	if len(hours) != NumHours {
+		t.Fatalf("len(Hours()) = %d, want %d", len(hours), NumHours)
+	}
+	if hours[0] != 0 {
+		t.Errorf("Hours()[0] = %d, want 0", hours[0])
+	}
+	if hours[len(hours)-1] != MaxHour {
+		t.Errorf("Hours()[last] = %d, want %d", hours[len(hours)-1], MaxHour)
+	}
+}
+
+func TestVariableIndex(t *testing.T) {
+	if VariableIndex(3, 5) != VarHeight {
+		t.Errorf("HGT: got %d, want %d", VariableIndex(3, 5), VarHeight)
+	}
+	if VariableIndex(2, 2) != VarWindU {
+		t.Errorf("UGRD: got %d, want %d", VariableIndex(2, 2), VarWindU)
+	}
+	if VariableIndex(2, 3) != VarWindV {
+		t.Errorf("VGRD: got %d, want %d", VariableIndex(2, 3), VarWindV)
+	}
+	if VariableIndex(0, 0) != -1 {
+		t.Errorf("unknown: got %d, want -1", VariableIndex(0, 0))
+	}
+}