diff --git a/model/config_test.go b/model/config_test.go new file mode 100644 index 0000000..bce288b --- /dev/null +++ b/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() +} diff --git a/service/rpc/io_stream_test.go b/service/rpc/io_stream_test.go new file mode 100644 index 0000000..ea90387 --- /dev/null +++ b/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} +}