code: def test_when_set_and_get_causal_model_then_the_set_model_is_returned():
      causal_dag = nx.DiGraph()
      causal_dag.add_node("X0")
      causal_model = ProbabilisticCausalModel(causal_dag)
  
      mdl = EmpiricalDistribution()
  
      causal_model.set_causal_mechanism("X0", mdl)
  
      assert causal_model.causal_mechanism("X0") == mdl
  
file_path: dowhy/tests/gcm/test_graph.py
imports: import networkx
  import numpy
  import pandas
  import pytest
  from flaky import flaky
  from pytest import approx
  from dowhy.gcm import AdditiveNoiseModel, EmpiricalDistribution, ProbabilisticCausalModel, draw_samples, fit
  from dowhy.gcm.ml import create_linear_regressor
  from dowhy.graph import is_root_node
label_desc: Use image node_id 4 for calling a global function with example usage: test_when_set_and_get_causal_model_then_the_set_model_is_returned() without return types