add tests for config & iostream (#1045)

* add tests for config & iostream

* chore

model/config_test.go

@@ -0,0 +1,169 @@
+package model
+
+import (
+	"os"
+	"strings"
+	"testing"
+)
+
+func TestReadConfig(t *testing.T) {
+	t.Run("ReadEmptyConfig", func(t *testing.T) {
+		file := newTempConfig(t, "")
+		c := &Config{}
+
+		if err := c.Read(file, nil); err != nil {
+			t.Fatalf("read empty config failed: %v", err)
+		}
+
+		testFields := []struct {
+			Name  string
+			Value any
+			Cond  bool
+		}{
+			{"jwt_secret_key", c.JWTSecretKey, c.JWTSecretKey != ""},
+			{"user_template", c.UserTemplate, c.UserTemplate == "user-dist"},
+			{"admin_template", c.AdminTemplate, c.AdminTemplate == "admin-dist"},
+			{"agent_secret_key", c.AgentSecretKey, c.AgentSecretKey != ""},
+		}
+
+		for _, field := range testFields {
+			if !field.Cond {
+				t.Fatalf("%s did not passed check, value: %v", field.Name, field.Value)
+			}
+		}
+
+		os.Remove(file)
+	})
+
+	t.Run("ReadFile", func(t *testing.T) {
+		const testCfg = "jwt_secret_key: test\nuser_template: um\nadmin_template: am\nagent_secret_key: none\nsite_name: lowkick"
+
+		var testFrontendTemplates = []FrontendTemplate{
+			{Path: "um"},
+			{Path: "am", IsAdmin: true},
+		}
+		file := newTempConfig(t, testCfg)
+		c := &Config{}
+
+		if err := c.Read(file, testFrontendTemplates); err != nil {
+			t.Fatalf("read config failed: %v", err)
+		}
+
+		testFields := []struct {
+			Name  string
+			Value any
+			Cond  bool
+		}{
+			{"jwt_secret_key", c.JWTSecretKey, c.JWTSecretKey == "test"},
+			{"user_template", c.UserTemplate, c.UserTemplate == "um"},
+			{"admin_template", c.AdminTemplate, c.AdminTemplate == "am"},
+			{"agent_secret_key", c.AgentSecretKey, c.AgentSecretKey == "none"},
+			{"site_name", c.SiteName, c.SiteName == "lowkick"},
+		}
+
+		for _, field := range testFields {
+			if !field.Cond {
+				t.Fatalf("%s did not passed check, value: %v", field.Name, field.Value)
+			}
+		}
+
+		os.Remove(file)
+	})
+
+	t.Run("ReadEnv", func(t *testing.T) {
+		os.Setenv("NZ_JWTSECRETKEY", "test")
+		os.Setenv("NZ_USERTEMPLATE", "um")
+		os.Setenv("NZ_ADMINTEMPLATE", "am")
+		os.Setenv("NZ_AGENTSECRETKEY", "none")
+		os.Setenv("NZ_HTTPS_LISTENPORT", "9876")
+
+		var testFrontendTemplates = []FrontendTemplate{
+			{Path: "um"},
+			{Path: "am", IsAdmin: true},
+		}
+		file := newTempConfig(t, "")
+		c := &Config{}
+
+		if err := c.Read(file, testFrontendTemplates); err != nil {
+			t.Fatalf("read empty config failed: %v", err)
+		}
+
+		testFields := []struct {
+			Name  string
+			Value any
+			Cond  bool
+		}{
+			{"jwt_secret_key", c.JWTSecretKey, c.JWTSecretKey == "test"},
+			{"user_template", c.UserTemplate, c.UserTemplate == "um"},
+			{"admin_template", c.AdminTemplate, c.AdminTemplate == "am"},
+			{"agent_secret_key", c.AgentSecretKey, c.AgentSecretKey == "none"},
+			{"https.listenport", c.HTTPS.ListenPort, c.HTTPS.ListenPort == 9876},
+		}
+
+		for _, field := range testFields {
+			if !field.Cond {
+				t.Fatalf("%s did not passed check, value: %v", field.Name, field.Value)
+			}
+		}
+
+		os.Remove(file)
+	})
+
+	t.Run("ReadEnvFile", func(t *testing.T) {
+		os.Setenv("NZ_JWTSECRETKEY", "test1")
+		os.Setenv("NZ_USERTEMPLATE", "um1")
+		os.Setenv("NZ_ADMINTEMPLATE", "am1")
+		os.Setenv("NZ_AGENTSECRETKEY", "none1")
+		os.Setenv("NZ_SITENAME", "lowkick1")
+
+		const testCfg = "jwt_secret_key: test\nuser_template: um\nadmin_template: am\nagent_secret_key: none\nsite_name: lowkick"
+
+		var testFrontendTemplates = []FrontendTemplate{
+			{Path: "um"},
+			{Path: "am", IsAdmin: true},
+		}
+		file := newTempConfig(t, testCfg)
+		c := &Config{}
+
+		if err := c.Read(file, testFrontendTemplates); err != nil {
+			t.Fatalf("read empty config failed: %v", err)
+		}
+
+		testFields := []struct {
+			Name  string
+			Value any
+			Cond  bool
+		}{
+			{"jwt_secret_key", c.JWTSecretKey, c.JWTSecretKey == "test"},
+			{"user_template", c.UserTemplate, c.UserTemplate == "um"},
+			{"admin_template", c.AdminTemplate, c.AdminTemplate == "am"},
+			{"agent_secret_key", c.AgentSecretKey, c.AgentSecretKey == "none"},
+			{"site_name", c.SiteName, c.SiteName == "lowkick"},
+		}
+
+		for _, field := range testFields {
+			if !field.Cond {
+				t.Fatalf("%s did not passed check, value: %v", field.Name, field.Value)
+			}
+		}
+
+		os.Remove(file)
+	})
+}
+
+func newTempConfig(t *testing.T, cfg string) string {
+	t.Helper()
+
+	file, err := os.CreateTemp(os.TempDir(), "nezha-test-config-*.yml")
+	if err != nil {
+		t.Fatalf("create temp file failed: %v", err)
+	}
+	defer file.Close()
+
+	_, err = file.ReadFrom(strings.NewReader(cfg))
+	if err != nil {
+		t.Fatalf("write to temp file failed: %v", err)
+	}
+
+	return file.Name()
+}

service/rpc/io_stream_test.go

@@ -0,0 +1,107 @@
+package rpc
+
+import (
+	"io"
+	"reflect"
+	"testing"
+	"time"
+)
+
+func TestIOStream(t *testing.T) {
+	handler := NewNezhaHandler()
+
+	const testStreamID = "ffffffff-ffff-ffff-ffff-ffffffffffff"
+
+	handler.CreateStream(testStreamID)
+	userIo, agentIo := newPipeReadWriter(), newPipeReadWriter()
+	defer func() {
+		userIo.Close()
+		agentIo.Close()
+	}()
+
+	handler.AgentConnected(testStreamID, agentIo)
+	handler.UserConnected(testStreamID, userIo)
+
+	go handler.StartStream(testStreamID, time.Second*10)
+
+	cases := [][]byte{
+		{0, 9, 1, 3, 2, 9, 1, 4, 8},
+		{3, 1, 3, 5, 2, 9, 5, 13, 53, 23},
+		make([]byte, 1024),
+		make([]byte, 1024*1024),
+	}
+
+	t.Run("WriteUserIO", func(t *testing.T) {
+		for i, c := range cases {
+			_, err := userIo.Write(c)
+			if err != nil {
+				t.Fatalf("write to userIo failed at case %d: %v", i, err)
+			}
+
+			b := make([]byte, len(c))
+			n, err := agentIo.Read(b)
+			if err != nil {
+				t.Fatalf("read agentIo failed at case %d: %v", i, err)
+			}
+
+			if !reflect.DeepEqual(c, b[:n]) {
+				t.Fatalf("expected %v, but got %v", c, b[:n])
+			}
+		}
+	})
+
+	t.Run("WriteAgentIO", func(t *testing.T) {
+		for i, c := range cases {
+			_, err := agentIo.Write(c)
+			if err != nil {
+				t.Fatalf("write to agentIo failed at case %d: %v", i, err)
+			}
+
+			b := make([]byte, len(c))
+			n, err := userIo.Read(b)
+			if err != nil {
+				t.Fatalf("read userIo failed at case %d: %v", i, err)
+			}
+
+			if !reflect.DeepEqual(c, b[:n]) {
+				t.Fatalf("Expected %v, but got %v", c, b[:n])
+			}
+		}
+	})
+
+	t.Run("WriteUserIOReadTwice", func(t *testing.T) {
+		data := []byte{1, 2, 3, 4, 5, 6, 7, 8}
+		_, err := agentIo.Write(data)
+		if err != nil {
+			t.Fatalf("write to agentIo failed: %v", err)
+		}
+
+		b := make([]byte, len(data)/2)
+		n, err := userIo.Read(b)
+		if err != nil {
+			t.Fatalf("read userIo failed: %v", err)
+		}
+
+		b2 := make([]byte, len(data)-n)
+		_, err = userIo.Read(b2)
+		if err != nil {
+			t.Fatalf("read userIo failed: %v", err)
+		}
+
+		if !reflect.DeepEqual(data[:len(data)/2], b) {
+			t.Fatalf("expected %v, but got %v", data[:len(data)/2], b)
+		}
+
+		if !reflect.DeepEqual(data[len(data)/2:], b2) {
+			t.Fatalf("expected %v, but got %v", data[len(data)/2:], b2)
+		}
+	})
+}
+
+func newPipeReadWriter() io.ReadWriteCloser {
+	r, w := io.Pipe()
+	return struct {
+		io.Reader
+		io.WriteCloser
+	}{r, w}
+}