func TestFileCustomRule(t *testing.T) { tqsc := tabletservermock.NewController() var qrs *tabletserver.QueryRules rulepath := path.Join(os.TempDir(), ".customrule.json") // Set r1 and try to get it back err := ioutil.WriteFile(rulepath, []byte(customRule1), os.FileMode(0644)) if err != nil { t.Fatalf("Cannot write r1 to rule file %s, err=%v", rulepath, err) } fcr := NewFileCustomRule() // Let FileCustomRule to build rule from the local file err = fcr.Open(tqsc, rulepath) if err != nil { t.Fatalf("Cannot open file custom rule service, err=%v", err) } // Fetch query rules we built to verify correctness qrs, _, err = fcr.GetRules() if err != nil { t.Fatalf("GetRules returns error: %v", err) } qr := qrs.Find("r1") if qr == nil { t.Fatalf("Expect custom rule r1 to be found, but got nothing, qrs=%v", qrs) } }
func TestZkCustomRule(t *testing.T) { tqsc := tabletservermock.NewController() setUpFakeZk(t) zkcr := NewZkCustomRule(conn) err := zkcr.Open(tqsc, "/zk/fake/customrules/testrules") if err != nil { t.Fatalf("Cannot open zookeeper custom rule service, err=%v", err) } var qrs *tabletserver.QueryRules // Test if we can successfully fetch the original rule (test GetRules) qrs, _, err = zkcr.GetRules() if err != nil { t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err) } qr := qrs.Find("r1") if qr == nil { t.Fatalf("Expect custom rule r1 to be found, but got nothing, qrs=%v", qrs) } // Test updating rules conn.Set("/zk/fake/customrules/testrules", customRule2, -1) <-time.After(time.Second) //Wait for the polling thread to respond qrs, _, err = zkcr.GetRules() if err != nil { t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err) } qr = qrs.Find("r2") if qr == nil { t.Fatalf("Expect custom rule r2 to be found, but got nothing, qrs=%v", qrs) } qr = qrs.Find("r1") if qr != nil { t.Fatalf("Custom rule r1 should not be found after r2 is set") } // Test rule path removal conn.Delete("/zk/fake/customrules/testrules", -1) <-time.After(time.Second) qrs, _, err = zkcr.GetRules() if err != nil { t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err) } if reflect.DeepEqual(qrs, tabletserver.NewQueryRules()) { t.Fatalf("Expect empty rule at this point") } // Test rule path revival conn.Create("/zk/fake/customrules/testrules", "customrule2", 0, zookeeper.WorldACL(zookeeper.PERM_ALL)) conn.Set("/zk/fake/customrules/testrules", customRule2, -1) <-time.After(time.Second) //Wait for the polling thread to respond qrs, _, err = zkcr.GetRules() if err != nil { t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err) } qr = qrs.Find("r2") if qr == nil { t.Fatalf("Expect custom rule r2 to be found, but got nothing, qrs=%v", qrs) } zkcr.Close() }