diff --git a/programs/test/benchmark.c b/programs/test/benchmark.c
index a7a0107..ba2c937 100755
--- a/programs/test/benchmark.c
+++ b/programs/test/benchmark.c
@@ -988,7 +988,7 @@
     }
 #endif
 
-#if defined(MBEDTLS_ECDH_C) && !defined(MBEDTLS_ECDH_LEGACY_CONTEXT)
+#if defined(MBEDTLS_ECDH_C) 
     if( todo.ecdh )
     {
         mbedtls_ecdh_context ecdh_srv, ecdh_cli;
@@ -1006,11 +1006,18 @@
             mbedtls_ecdh_setup( &ecdh_srv, curve_info->grp_id );
             mbedtls_ecdh_setup( &ecdh_cli, curve_info->grp_id );
 
+#if defined(MBEDTLS_ECDH_LEGACY_CONTEXT)
+            if (mbedtls_ecp_group_load(&ecdh_srv.grp, curve_info->grp_id) != 0 ||
+                mbedtls_ecdh_gen_public(&ecdh_srv.grp,
+                    &ecdh_srv.d,
+                    &ecdh_srv.Q, myrand, NULL) != 0)
+#else
             if( ecdh_srv.var == MBEDTLS_ECDH_VARIANT_MBEDTLS_2_0 && (
                 mbedtls_ecp_group_load( &ecdh_srv.ctx.mbed_ecdh.grp, curve_info->grp_id ) != 0 ||
                 mbedtls_ecdh_gen_public( &ecdh_srv.ctx.mbed_ecdh.grp,
                                          &ecdh_srv.ctx.mbed_ecdh.d,
                                          &ecdh_srv.ctx.mbed_ecdh.Q, myrand, NULL ) != 0 ))
+#endif
                 mbedtls_exit( 1 );
 
             mbedtls_snprintf( title, sizeof( title ), "ECDHE-%s", curve_info->name );