diff --git a/ctx/go-http-context/context_test.go b/ctx/go-http-context/context_test.go index 65a1936..9a4abe4 100644 --- a/ctx/go-http-context/context_test.go +++ b/ctx/go-http-context/context_test.go @@ -17,7 +17,7 @@ func TestHttpContext_IsValidContext(t *testing.T) { } request, err := http.NewRequest("GET", "https://baidu.com", strings.NewReader("")) if err != nil { - + t.Error(err) } tests := []struct { name string @@ -126,7 +126,7 @@ func TestHttpReqStorage_Delete(t *testing.T) { r := HttpReqStorage{ source: tt.fields.source, } - r.Delete(tt.args.key) + r.Delete(ctx.StorageKey(tt.args.key)) }) } } @@ -151,7 +151,7 @@ func TestHttpReqStorage_Get(t *testing.T) { r := HttpReqStorage{ source: tt.fields.source, } - if got := r.Get(tt.args.key); !reflect.DeepEqual(got, tt.want) { + if got := r.Get(ctx.StorageKey(tt.args.key)); !reflect.DeepEqual(got, tt.want) { t.Errorf("Get() = %v, want %v", got, tt.want) } }) @@ -178,7 +178,7 @@ func TestHttpReqStorage_Set(t *testing.T) { r := HttpReqStorage{ source: tt.fields.source, } - r.Set(tt.args.key, tt.args.value) + r.Set(ctx.StorageKey(tt.args.key), tt.args.value) }) } } diff --git a/ctx/go-http-context/request-storage.go b/ctx/go-http-context/request-storage.go index f45fa29..3fa55fd 100644 --- a/ctx/go-http-context/request-storage.go +++ b/ctx/go-http-context/request-storage.go @@ -2,6 +2,7 @@ package go_http_context import ( "context" + "github.com/weloe/token-go/ctx" "net/http" ) @@ -13,18 +14,18 @@ func NewReqStorage(req *http.Request) *HttpReqStorage { return &HttpReqStorage{source: req.Context()} } -func (r HttpReqStorage) Source() interface{} { +func (r *HttpReqStorage) Source() interface{} { return r.source } -func (r HttpReqStorage) Get(key string) interface{} { +func (r *HttpReqStorage) Get(key ctx.StorageKey) interface{} { return r.source.Value(key) } -func (r HttpReqStorage) Set(key string, value string) { +func (r *HttpReqStorage) Set(key ctx.StorageKey, value string) { r.source = context.WithValue(r.source, key, value) } -func (r HttpReqStorage) Delete(key string) { +func (r *HttpReqStorage) Delete(key ctx.StorageKey) { r.source = context.WithValue(r.source, key, nil) } diff --git a/ctx/request_storage.go b/ctx/request_storage.go index 7f473e5..d5f6ff8 100644 --- a/ctx/request_storage.go +++ b/ctx/request_storage.go @@ -1,10 +1,12 @@ package ctx +type StorageKey string + // ReqStorage // Use to set-get-delete data,and it'll be cleaned after request type ReqStorage interface { Source() interface{} - Get(key string) interface{} - Set(key string, value string) - Delete(key string) + Get(key StorageKey) interface{} + Set(key StorageKey, value string) + Delete(key StorageKey) }