Support setting amount to be generated

4 files changed, 38 insertions, 39 deletions

diff --git a/lib/ast_types.ml b/lib/ast_types.ml
index 7b4f445..b9f4294 100644
--- a/lib/ast_types.ml
+++ b/lib/ast_types.ml
@@ -11,7 +11,7 @@ type ast_row =
 [@@deriving show, eq]
 
 type ast_table =
-  | Table of string * ast_row * ast_table
+  | Table of string * int * ast_row * ast_table
   | End
 [@@deriving show, eq]
 
@@ -32,8 +32,8 @@ let rec print_row = function
 
 
 let rec print = function
-  | Table (s, r, t) ->
-    printf "%s:" s;
+  | Table (s, amount, r, t) ->
+    printf "%s (%d):" s amount;
     print_row r;
     print t
   | End -> printf "\n"
@@ -46,13 +46,13 @@ type ast_row_relation =
   | CRow of (string * string * (string list -> Fixture.t))
 
 type ast_table_relation =
-  | PTable of string * Fixture.t list
-  | CTable of string * string * string * (string list -> Fixture.t list)
+  | PTable of string * int *  Fixture.t list
+  | CTable of string * string * string * int * (string list -> Fixture.t list)
   (** parent, row, child_name, row *)
 
 let table_name = function 
-  | PTable (n,_) -> n 
-  | CTable (_,_,c,_) -> c
+  | PTable (n,_,_) -> n 
+  | CTable (_,_,c,_,_) -> c
 
 let rec ast_row_to_fixtures = function
   | Uuidv4 (s, r) -> PRow (Fixture.Uuidv4 s) :: ast_row_to_fixtures r
@@ -63,7 +63,7 @@ let rec ast_row_to_fixtures = function
   | End -> []
 
 
-let ast_table_to_table name rows =
+let ast_table_to_table name amount rows  =
   (** Convert an AST table to a Fixture table that we can more easily work with for
      feeding the generated data as input into fixtures that depend on the generated
      data from other fixtures *)
@@ -72,13 +72,13 @@ let ast_table_to_table name rows =
       match row with
       | PRow r ->
         (match tbl with
-         | PTable (n, l) -> PTable (n, r :: l)
-         | CTable (p, pr, c, f) -> CTable (p, pr, c, fun l -> r :: f l))
+         | PTable (n, amount, l) -> PTable (n, amount, r :: l)
+         | CTable (p, pr, c, amount, f) -> CTable (p, pr, c, amount, fun l -> r :: f l))
       | CRow (p, pr, f) ->
         (match tbl with
-         | PTable (c, lr) -> CTable (p, pr, c, fun l -> f l :: lr)
+         | PTable (c,amount, lr ) -> CTable (p, pr, c, amount, fun l -> f l :: lr)
          | CTable _ -> failwith "Cannot have multiple relations"))
-    (PTable (name, []))
+    (PTable (name, amount, []))
     rows
 
 
@@ -86,20 +86,20 @@ let rec ast_table_to_tables ast =
   (** Convert an AST table to a list of Fixture tables, so once data has been generated
       it can be used as input for the next entry *)
   match ast with
-  | Table (name, r, t) ->
+  | Table (name, amount, r, t) ->
     print_endline ("Converting table: " ^ name);
-    ast_table_to_table name (ast_row_to_fixtures r) :: ast_table_to_tables t
+    ast_table_to_table name amount (ast_row_to_fixtures r) :: ast_table_to_tables t
   | End -> []
 
 let%test "ast_table_to_tables" = 
-  let purchase = Table ("purchase", Uuidv4 ("id", Name ("name", End)), End) in
-  let ast = Table ("user", Uuidv4 ("id", Name ("name", End)), purchase) in
+  let purchase = Table ("purchase", 5, Uuidv4 ("id", Name ("name", End)), End) in
+  let ast = Table ("user", 2, Uuidv4 ("id", Name ("name", End)), purchase) in
   let tables = ast_table_to_tables ast in 
   List.length tables == 2
 
 let generated_fixtures = ref (Hashtbl.create 10)
 
-let rec resolve_fixtures tables name ~amount = 
+let rec resolve_fixtures tables name = 
   (* Verify if we have already generated fixtures for this table *) 
   let maybe_fixtures = Hashtbl.find_opt !generated_fixtures name in 
   match maybe_fixtures with 
@@ -108,13 +108,13 @@ let rec resolve_fixtures tables name ~amount =
   let maybe_table = List.find_opt (fun tbl -> table_name tbl = name) tables 
   in 
   match maybe_table with 
-  | Some (PTable (n,l)) -> 
+  | Some (PTable (n, amount, l)) -> 
       let resolved = (n, Fixture.compile l ~amount) in 
       (* store the generated fixtures for this table *)
       Hashtbl.add !generated_fixtures name (snd resolved);
       resolved 
-  | Some (CTable (p,pr,c,create_fixtures)) ->  (
-      let (_, parent_fixtures) = resolve_fixtures tables p ~amount in
+  | Some (CTable (p,pr,c, amount, create_fixtures)) ->  (
+      let (_, parent_fixtures) = resolve_fixtures tables p in
       (* store the generated fixtures for this table *) 
       Hashtbl.add !generated_fixtures name parent_fixtures;
       print_endline ("Resolving fixtures for " ^ c);
@@ -131,11 +131,11 @@ let rec resolve_fixtures tables name ~amount =
 
 let%test "resolve_fixtures" = 
   let tables = [
-    PTable ("user", [Fixture.Uuidv4 "id"; Fixture.Uuidv4 "uuid"]);
-    CTable ("user", "id", "posts", fun ids -> [Fixture.Foreign ("user_id", ids)])
+    PTable ("user", 1, [Fixture.Uuidv4 "id"; Fixture.Uuidv4 "uuid"]);
+    CTable ("user", "id", "posts", 1, fun ids -> [Fixture.Foreign ("user_id", ids)])
   ] in 
-  let (_, user_fixtures) = resolve_fixtures tables "user" ~amount:1 in 
-  let (_, purchase_fixtures) = resolve_fixtures tables "posts" ~amount:1 in 
+  let (_, user_fixtures) = resolve_fixtures tables "user"  in 
+  let (_, purchase_fixtures) = resolve_fixtures tables "posts" in 
         for i = 0 to List.length user_fixtures - 1 do 
           print_endline ("User fixture: " ^ Fixture.csv_of_string_list (List.nth user_fixtures i));
         done;
@@ -150,24 +150,24 @@ let%test "resolve_fixtures" =
 
 let show_tables tables = 
   let show_table = function 
-    | PTable (n,_) -> n 
-    | CTable (_,_,c,_) -> c 
+    | PTable (n,_,_) -> n 
+    | CTable (_,_,c,_,_) -> c 
   in 
   tables |> List.map show_table |> String.concat "\n"
 
 
-let compile ast ~amount =
+let compile ast =
   let tables = ast_table_to_tables ast in
   print_endline ("Table length: " ^ string_of_int (List.length tables));
   print_endline (show_tables tables);
   let resolve_fixtures' = function 
-    | PTable (n,l) -> 
-        let resolved = (n, Fixture.compile l ~amount) in 
+    | PTable (n, amount, l) -> 
+        let (n,_,resolved) = (n, amount, Fixture.compile l ~amount) in 
         (* store the generated fixtures for this table *)
-        Hashtbl.add !generated_fixtures n (snd resolved);
-        resolved
-    | CTable (p,pr,c,create_fixtures) -> 
-      let (_, parent_fixtures) = resolve_fixtures tables p ~amount in
+        Hashtbl.add !generated_fixtures n resolved;
+        (n,resolved)
+    | CTable (p,pr,c, amount, create_fixtures) -> 
+      let (_, parent_fixtures) = resolve_fixtures tables p in
       (* store the generated fixtures for this table *)
       Hashtbl.add !generated_fixtures c parent_fixtures;
       print_endline ("Resolving fixtures for " ^ c);
diff --git a/lib/fixture.ml b/lib/fixture.ml
index 8806692..278024c 100644
--- a/lib/fixture.ml
+++ b/lib/fixture.ml
@@ -22,7 +22,6 @@ let names =
 let random_name () = List.nth names (Random.int (List.length names))
 let%test "random_name" = List.mem (random_name ()) names
 
-(* TODO: Support arbitrary amount *)
 (* TODO: Support names *)
 type t =
   | Name of string
diff --git a/lib/lexer.mll b/lib/lexer.mll
index fed77ab..e0c5453 100644
--- a/lib/lexer.mll
+++ b/lib/lexer.mll
@@ -21,8 +21,8 @@ rule read =
   | "_uuidv4" { UUIDV4 }
   | "_name" { NAME }
   | "_int" { INTSYMBOL } 
-  | "<" { LBRACE }
-  | ">" { RBRACE }
+  | "(" { LBRACE }
+  | ")" { RBRACE }
   | int { INT (Lexing.lexeme lexbuf |> int_of_string) }
   | white { read lexbuf } 
   | ":" { COLON }
diff --git a/lib/parser.mly b/lib/parser.mly
index ab9b697..e9d06bc 100644
--- a/lib/parser.mly
+++ b/lib/parser.mly
@@ -22,9 +22,9 @@ prog:
   ;
 
 expr:
-  | tbl = IDENTIFIER; COLON; r = row; NEWLINE; e = expr { Table (tbl,r, e) }
-  | tbl = IDENTIFIER; COLON; r = row; SEMICOLON; e = expr { Table (tbl,r, e) }
-  | tbl = IDENTIFIER; COLON; r = row { Table (tbl,r, End) }
+  | tbl = IDENTIFIER; LBRACE; amount = INT; RBRACE; COLON; r = row; NEWLINE; e = expr { Table (tbl, amount, r, e) }
+  | tbl = IDENTIFIER; LBRACE; amount = INT; RBRACE; COLON; r = row; SEMICOLON; e = expr { Table (tbl,amount, r, e) }
+  | tbl = IDENTIFIER; LBRACE; amount = INT; RBRACE; COLON; r = row { Table (tbl,amount, r, End) }
   ;
 
 row: