diff --git a/lab-4/main_test.c b/lab-4/main_test.c
index 0e80882..057bed9 100644
--- a/lab-4/main_test.c
+++ b/lab-4/main_test.c
@@ -43,9 +43,9 @@ UTEST(parser, missing_semi_colon) {
   ASSERT_EQ(result, 1); 
 }
 
-UTEST(parser, sample) {
+UTEST(parser, hello) {
   // Read sample file as input
-  yyin = fopen("samples/program.c", "r");
+  yyin = fopen("samples/hello-world.cbl", "r");
   yyrestart(yyin);
   ASSERT_TRUE(yyin);
 
@@ -55,3 +55,31 @@ UTEST(parser, sample) {
   // Assert the result to test correctness
   ASSERT_EQ(result, 0); 
 }
+
+UTEST(parser, print) {
+  // Must include the null character to terminate input
+  char string[] = "DISPLAY 'Hello World!'\0"; 
+  YY_BUFFER_STATE buffer = yy_scan_buffer(string, sizeof(string));
+
+  yylineno = 1;
+  int result = yyparse();
+
+  yy_delete_buffer(buffer);
+
+  // Assert the result to test correctness
+  ASSERT_EQ(result, 0); 
+}
+
+UTEST(parser, branching) {
+  // Must include the null character to terminate input
+  char string[] = "IF A > B THEN DISPLAY 'A is greater than B' ELSE DISPLAY 'B is greater than A'\0"; 
+  YY_BUFFER_STATE buffer = yy_scan_buffer(string, sizeof(string));
+
+  yylineno = 1;
+  int result = yyparse();
+
+  yy_delete_buffer(buffer);
+
+  // Assert the result to test correctness
+  ASSERT_EQ(result, 0); 
+}  
\ No newline at end of file
diff --git a/lab-4/parser.bison b/lab-4/parser.bison
index ddbc8b4..e4c4b6f 100644
--- a/lab-4/parser.bison
+++ b/lab-4/parser.bison
@@ -77,6 +77,7 @@ sect_data       : TOKEN_PROGRAM_ID  TOKEN_DOT TOKEN_IDENT TOKEN_DOT
 type            : TOKEN_IDENTIFICATION
                 | TOKEN_PROCEDURE
                 | TOKEN_STOP
+                | TOKEN_KEYWORD_DATA
                 ;
 simple_stmt     : function
                 ;