/*
 *      Copyright (C) 2012-2015 DataStax Inc.
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */
package com.datastax.driver.core;

import com.datastax.driver.core.policies.ConstantSpeculativeExecutionPolicy;
import com.datastax.driver.core.utils.CassandraVersion;
import org.mockito.Mockito;
import org.testng.annotations.Test;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.times;

/**
 * Test that needs a real CCM cluster (as opposed to SCassandra for other specex tests), because
 * it uses a protocol v3 feature.
 */
public class SpeculativeExecutionIntegrationTest extends CCMTestsSupport {

    TimestampGenerator timestampGenerator;

    @Override
    public void onTestContextInitialized() {
        execute("create table foo(k int primary key, v int)");
    }

    @Override
    public Cluster.Builder createClusterBuilder() {
        timestampGenerator = Mockito.spy(ServerSideTimestampGenerator.INSTANCE);
        return Cluster.builder()
                .withTimestampGenerator(timestampGenerator)
                .withQueryOptions(new QueryOptions().setDefaultIdempotence(true))
                        // Set an artificially low timeout to force speculative execution
                .withSpeculativeExecutionPolicy(new ConstantSpeculativeExecutionPolicy(1, 2));
    }

    /**
     * Validates that if a query gets speculatively re-executed, the second execution uses the same default timestamp.
     *
     * @test_category tracing
     * @jira_ticket JAVA-724
     * @expected_result timestamp generator invoked only once for a query that caused two executions.
     */
    @Test(groups = "short")
    @CassandraVersion(major = 2.1)
    public void should_use_same_default_timestamp_for_all_executions() {
        Metrics.Errors errors = cluster().getMetrics().getErrorMetrics();

        Mockito.reset(timestampGenerator);
        long execStartCount = errors.getSpeculativeExecutions().getCount();

        BatchStatement batch = new BatchStatement();
        for (int k = 0; k < 1000; k++) {
            batch.add(new SimpleStatement("insert into foo(k,v) values (?,1)", k));
        }
        session().execute(batch);

        assertThat(errors.getSpeculativeExecutions().getCount()).isEqualTo(execStartCount + 1);
        Mockito.verify(timestampGenerator, times(1)).next();
    }
}

@@ -900,26 +900,20 @@
 ertions.
-assertThat
+fail
 ;%0Aimport
@@ -1580,98 +1580,8 @@
 or)%0A
-                .withQueryOptions(new QueryOptions().setDefaultIdempotence(true))%0A        
     
@@ -2270,24 +2270,411 @@
 Metrics();%0A%0A
+        // The check is attempted up to 10 times to account for the small possibility that a%0A        // scheduled execution is not needed/exercised.  Even though the policy is set up%0A        // to schedule an execution after 1ms, the timeout might not fire before the response is received.%0A        int tryCount = 0;%0A        int maxTries = 10;%0A        while (tryCount++ %3C maxTries) %7B%0A    
         Mock
@@ -2704,16 +2704,20 @@
 rator);%0A
+    
         
@@ -2793,16 +2793,20 @@
         
+   
+ 
 BatchSta
@@ -2838,24 +2838,28 @@
 tatement();%0A
+    
         for 
@@ -2899,16 +2899,20 @@
         
+    
 batch.ad
@@ -2972,16 +2972,36 @@
 ,1)%22, k)
+.setIdempotent(true)
 );%0A     
@@ -3003,18 +3003,65 @@
         
+  
+  
 %7D%0A
+            batch.setIdempotent(true);%0A    
         
@@ -3099,18 +3099,15 @@
     
-assertThat
+    if 
 (err
@@ -3151,20 +3151,12 @@
 nt()
-).isEqualTo(
+ == 
 exec
@@ -3170,26 +3170,31 @@
 unt + 1)
-;%0A
+ %7B%0A    
         
 Mockito.
@@ -3177,32 +3177,36 @@
 ) %7B%0A            
+    
 Mockito.verify(t
@@ -3238,20 +3238,185 @@
 1)).next();%0A
+                break;%0A            %7D%0A        %7D%0A%0A        if (tryCount == maxTries) %7B%0A            fail(%22Observed no speculative executions in 10 attempts%22);%0A        %7D%0A
     %7D%0A%7D%0A

Set idempotence on Statement instead of globally