77from skyflow .errors ._skyflow_errors import SkyflowError , SkyflowErrorCodes , SkyflowErrorMessages
88from skyflow .vault ._client import Client , Configuration
99from skyflow .service_account import generate_bearer_token
10+ from skyflow .vault ._config import RedactionType
1011from dotenv import dotenv_values
1112import warnings
1213
@@ -55,7 +56,8 @@ def testGetDetokenizeRequestBodyWithValidBody(self):
5556 body = getDetokenizeRequestBody (self .tokenField )
5657 expectedOutput = {
5758 "detokenizationParameters" : [{
58- "token" : self .testToken
59+ "token" : self .testToken ,
60+ "redaction" : "PLAIN_TEXT"
5961 }]
6062 }
6163
@@ -101,6 +103,15 @@ def testDetokenizeTokenInvalidType(self):
101103 self .assertEqual (
102104 e .message , SkyflowErrorMessages .INVALID_TOKEN_TYPE .value % (list ))
103105
106+ def testDetokenizeRedactionInvalidType (self ):
107+ invalidData = {"records" : [{"token" : "valid" , "redaction" : 'demo' }]}
108+ try :
109+ self .client .detokenize (invalidData )
110+ except SkyflowError as error :
111+ self .assertTrue (error )
112+ self .assertEqual (error .code , SkyflowErrorCodes .INVALID_INPUT .value )
113+ self .assertEqual (error .message , SkyflowErrorMessages .INVALID_REDACTION_TYPE .value % str (type ("demo" )))
114+
104115 def testResponseBodySuccess (self ):
105116 response = {"records" : [{"token" : "abc" , "value" : "secret" }]}
106117 self .add_mock_response (response , 200 )
@@ -135,3 +146,39 @@ def testResponseNotJson(self):
135146 self .assertEqual (error .code , 200 )
136147 self .assertEqual (error .message , expectedError .value %
137148 response .decode ('utf-8' ))
149+
150+ def testRequestBodyNoRedactionKey (self ):
151+ expectedOutput = {
152+ "detokenizationParameters" : [{
153+ "token" : self .testToken ,
154+ "redaction" : "PLAIN_TEXT"
155+ }]
156+ }
157+ requestBody = getDetokenizeRequestBody (self .tokenField )
158+ self .assertEqual (requestBody , expectedOutput )
159+
160+ def testRequestBodyWithValidRedaction (self ):
161+ expectedOutput = {
162+ "detokenizationParameters" : [{
163+ "token" : self .testToken ,
164+ "redaction" : "REDACTED"
165+ }]
166+ }
167+ data = {
168+ "token" : self .testToken ,
169+ "redaction" : RedactionType .REDACTED
170+ }
171+ requestBody = getDetokenizeRequestBody (data )
172+ self .assertEqual (expectedOutput , requestBody )
173+
174+ def testRequestBodyWithInValidRedaction (self ):
175+ data = {
176+ "token" : self .testToken ,
177+ "redaction" : "123"
178+ }
179+ try :
180+ getDetokenizeRequestBody (data )
181+ except SkyflowError as error :
182+ self .assertTrue (error )
183+ self .assertEqual (error .code , SkyflowErrorCodes .INVALID_INPUT .value )
184+ self .assertEqual (error .message , SkyflowErrorMessages .INVALID_REDACTION_TYPE .value % str (type (data ["redaction" ])))
0 commit comments