1 files changed, 5 insertions, 15 deletions
diff --git a/tests/wycheproof/test_ecdh.py b/tests/wycheproof/test_ecdh.py
index a0c3f0d9..5fcc45b7 100644
--- a/tests/wycheproof/test_ecdh.py
+++ b/tests/wycheproof/test_ecdh.py
@@ -97,28 +97,18 @@ def test_ecdh_ecpoint(backend, wycheproof):
     private_key = ec.derive_private_key(
         int(wycheproof.testcase["private"], 16), curve, backend
     )
-    # We don't support compressed points
-    if (
-        wycheproof.has_flag("CompressedPoint") or
-        not wycheproof.testcase["public"]
-    ):
+
+    if wycheproof.invalid:
         with pytest.raises(ValueError):
-            ec.EllipticCurvePublicNumbers.from_encoded_point(
+            ec.EllipticCurvePublicKey.from_encoded_point(
                 curve, binascii.unhexlify(wycheproof.testcase["public"])
             )
         return
 
-    public_numbers = ec.EllipticCurvePublicNumbers.from_encoded_point(
+    assert wycheproof.valid or wycheproof.acceptable
+    public_key = ec.EllipticCurvePublicKey.from_encoded_point(
         curve, binascii.unhexlify(wycheproof.testcase["public"])
     )
-    if wycheproof.testcase["comment"] == "point is not on curve":
-        assert wycheproof.invalid
-        with pytest.raises(ValueError):
-            public_numbers.public_key(backend)
-        return
-
-    assert wycheproof.valid or wycheproof.acceptable
-    public_key = public_numbers.public_key(backend)
     computed_shared = private_key.exchange(ec.ECDH(), public_key)
     expected_shared = binascii.unhexlify(wycheproof.testcase["shared"])
     assert computed_shared == expected_shared