import pytest
import networkx as nx
import numpy as np
from quantumwalks import GraphDTQW

class TestGraphDTQW:
    def setup_method(self):
        self.graph = nx.cycle_graph(5)
    
    def test_initialization(self):
        gqw = GraphDTQW(self.graph, coin_type='grover')
        assert gqw.num_nodes == 5
        assert np.isclose(np.linalg.norm(gqw.initial_state), 1.0)
    
    def test_probability_conservation(self):
        gqw = GraphDTQW(self.graph, coin_type='grover')
        gqw.evolve(5)
        probs = gqw.get_probabilities()
        assert np.isclose(np.sum(probs), 1.0)
    
    def test_finite_average(self):
        gqw = GraphDTQW(self.graph, coin_type='grover')
        avg_probs = gqw.average_probabilities_finite(10)
        assert np.isclose(np.sum(avg_probs), 1.0)
        assert len(avg_probs) == gqw.num_nodes