Import libs 4 (#758)

1 files changed, 15 insertions, 2 deletions

diff --git a/contrib/python/google-auth/py3/tests/crypt/test_es256.py b/contrib/python/google-auth/py3/tests/crypt/test_es256.py
index 1a43a2f01b..3ba5b64fad 100644
--- a/contrib/python/google-auth/py3/tests/crypt/test_es256.py
+++ b/contrib/python/google-auth/py3/tests/crypt/test_es256.py
@@ -15,6 +15,7 @@
 import base64
 import json
 import os
+import pickle
 
 from cryptography.hazmat.primitives.asymmetric import ec
 import pytest  # type: ignore
@@ -24,8 +25,8 @@ from google.auth.crypt import base
 from google.auth.crypt import es256
 
 
-import yatest.common
-DATA_DIR = os.path.join(yatest.common.test_source_path(), "data")
+import yatest.common as yc
+DATA_DIR = os.path.join(os.path.dirname(yc.source_path(__file__)), "..", "data")
 
 # To generate es256_privatekey.pem, es256_privatekey.pub, and
 # es256_public_cert.pem:
@@ -142,3 +143,15 @@ class TestES256Signer(object):
 
         assert signer.key_id == SERVICE_ACCOUNT_INFO[base._JSON_FILE_PRIVATE_KEY_ID]
         assert isinstance(signer._key, ec.EllipticCurvePrivateKey)
+
+    def test_pickle(self):
+        signer = es256.ES256Signer.from_service_account_file(SERVICE_ACCOUNT_JSON_FILE)
+
+        assert signer.key_id == SERVICE_ACCOUNT_INFO[base._JSON_FILE_PRIVATE_KEY_ID]
+        assert isinstance(signer._key, ec.EllipticCurvePrivateKey)
+
+        pickled_signer = pickle.dumps(signer)
+        signer = pickle.loads(pickled_signer)
+
+        assert signer.key_id == SERVICE_ACCOUNT_INFO[base._JSON_FILE_PRIVATE_KEY_ID]
+        assert isinstance(signer._key, ec.EllipticCurvePrivateKey)